mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
feat(download): fill model metadata from the source API on download
A ModelScope or Hugging Face download landed as a bare filename, hash and
source link; the model card stayed empty until the user ran "Enrich
Metadata with AI" by hand. But everything that makes a CivitAI download
useful — the display name, the description, the tags, the trigger words,
the example images, the preview — is already published by those sites'
public APIs, so asking for it at download time is deterministic work, not
model work.
Add `py/services/model_sources/hydration.py`, called by
`_save_source_metadata()` once the sidecar exists and the file is in the
scanner cache. It fetches the model card plus the site's card extras and
hands them to the same `PostProcessor` the AI skill uses, with an empty
`llm_output`, so the two paths cannot drift apart. What lands:
* `model_name` from the site's own display name (ModelScope's `Name`), so
the card stops showing the local filename — written only while the value
still equals the file stem, since once a user renames a model that
choice is theirs to keep
* `civitai.name` from the matched version's label (`showName`), which the
card renders as the version chip
* `civitai.description` / `modelDescription` from the author summary plus
the README as HTML
* `civitai.images` / `preview_url` from the per-file example images
* `civitai.trainedWords` from the per-file trigger words
* `base_model`, `tags` and `usage_tips` as before
Provenance stays honest: the pass records
`metadata_source = "source:<platform>"` rather than the skill's
`agent:enrich_hf_metadata`, and — because no provider ran — it no longer
stamps `llm_enriched_at`; that stamp is now conditional on the LLM
actually answering, which is what the field means. The five hand-rolled
`civitai` dict merges in the post-processor collapse into one
`_merge_civitai()` helper.
Two guards keep it safe. Only a model whose stored
`source_platform`/`source_url` match the repository being downloaded is
updated, so a local file that merely shares a name never receives another
model's card; and a file already on disk is topped up too, which
back-fills models downloaded before this existed. READMEs and detail
payloads describe the repository rather than the file, so a short-lived
process-wide `ModelSourceCache` (300 s, 32 entries) keeps a batch over one
repository to two HTTP requests. Every failure is logged and swallowed:
hydration can never fail a download.
Fix the hash policy while here. `_save_source_metadata()` went straight to
`MetadataManager.create_default_metadata()`, bypassing the per-type
factory on the owning scanner, so a checkpoint paid a full SHA256 inside
the download request — `CheckpointScanner`/`OtherScanner` deliberately
record `hash_status="pending"` with an empty `sha256` for their multi-GB
files. Metadata is now created through `scanner._create_default_metadata()`.
Hydration copes with the empty hash: `_matching_versions()` falls back to
the repository basename, which is exactly what the download just wrote.
Report both post-transfer stages, which advance no byte counter and so
read as a stall: the bar sat at 100% showing `0 B/s` for the seconds spent
hashing and fetching. `_report_phase()` broadcasts
`{"status": "metadata", "stage": "indexing" | "source", "platform": ...}`,
and `LoadingManager` names the stage in the status line (keeping the batch
position), retitles the item line, replaces the dead speed figure and runs
a sheen over the bar. `stage`/`platform` are machine-readable; the wording
is localised in the frontend.
Finally, `modelscope.ai` is its own catalogue rather than an alias of
`modelscope.cn` — `referall13/EM1` exists only on `.ai` and
`jj3550945163/Krea-2-LORA` only on `.cn` — so its URLs were rejected with
"Invalid model URL format". Register it as `ModelScopeIntlSource`
(`platform="modelscope-ai"`, `msai:` group prefix, its own default
download directory) and derive every URL either deployment builds from a
per-class `base_url`. `modelscope.com` stays an alias of `.cn`, which is
what it redirects to. The frontend source table, the link dialog hints and
the docs mirror the split.
Verified against the live APIs: both reported `.ai` repositories list
their files, read their READMEs and yield name / version / base model /
trigger words / example images. Backend 3092 passed; frontend 1259 JS +
91 Vue passed. The nine locales carry the new progress copy in the next
commit.
This commit is contained in:
Will Miao
@@ -0,0 +1,397 @@
|
||||
"""Tests for download-time metadata hydration.
|
||||
|
||||
`py/services/model_sources/hydration.py` is the deterministic counterpart of
|
||||
the `enrich_hf_metadata` skill: it turns a freshly downloaded ModelScope /
|
||||
Hugging Face file into the populated model card a CivitAI download produces,
|
||||
without an LLM and without the user running anything.
|
||||
|
||||
These tests cover the orchestration — which source data is fetched, what is
|
||||
handed to the post-processor, and that nothing here can fail a download. The
|
||||
field-by-field mapping lives in `tests/services/test_post_processor.py`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.model_sources import ModelCardContext, ModelSourceCache, SourceRef
|
||||
from py.services.model_sources import hydration
|
||||
from py.services.model_sources.base import ModelSource
|
||||
from py.services.model_sources.hydration import (
|
||||
SHARED_CACHE_MAX_ENTRIES,
|
||||
hydrate_from_source,
|
||||
load_model_card,
|
||||
reset_shared_caches,
|
||||
resolve_site_base_model,
|
||||
shared_source_cache,
|
||||
)
|
||||
|
||||
REF = SourceRef(
|
||||
platform="modelscope",
|
||||
source_id="user/repo",
|
||||
url="https://modelscope.cn/models/user/repo",
|
||||
)
|
||||
|
||||
SIDECAR = {
|
||||
"sha256": "a" * 64,
|
||||
"base_model": "Unknown",
|
||||
# Written by the download handler just before hydration runs.
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/user/repo",
|
||||
}
|
||||
|
||||
|
||||
class _FakeSource(ModelSource):
|
||||
"""Minimal provider that records what hydration asked of it."""
|
||||
|
||||
platform = "modelscope"
|
||||
label = "ModelScope"
|
||||
supports_enrichment = True
|
||||
|
||||
def __init__(self, *, context=None, readme="", fail=False):
|
||||
self.context = context if context is not None else ModelCardContext()
|
||||
self.readme = readme
|
||||
self.fail = fail
|
||||
self.readme_calls = 0
|
||||
self.context_calls = 0
|
||||
self.context_kwargs: dict = {}
|
||||
|
||||
async def fetch_model_card(self, source_id):
|
||||
self.readme_calls += 1
|
||||
if self.fail:
|
||||
raise RuntimeError("network down")
|
||||
return self.readme
|
||||
|
||||
async def fetch_model_card_context(
|
||||
self, source_id, filename="", *, sha256="", cache=None
|
||||
):
|
||||
self.context_calls += 1
|
||||
self.context_kwargs = {"filename": filename, "sha256": sha256}
|
||||
if self.fail:
|
||||
raise RuntimeError("network down")
|
||||
return self.context
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolated_shared_caches():
|
||||
reset_shared_caches()
|
||||
yield
|
||||
reset_shared_caches()
|
||||
|
||||
|
||||
def _async(value):
|
||||
async def _call(*_args, **_kwargs):
|
||||
return value
|
||||
|
||||
return _call
|
||||
|
||||
|
||||
def _wire(monkeypatch, source, *, metadata=SIDECAR, result=None):
|
||||
"""Patch hydration's collaborators; return the recorded process() calls."""
|
||||
|
||||
monkeypatch.setattr(hydration, "get_source", lambda _platform: source)
|
||||
monkeypatch.setattr("py.metadata_ops.read_metadata", _async(metadata))
|
||||
|
||||
calls: list = []
|
||||
|
||||
class _Processor:
|
||||
async def process(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
if result is not None:
|
||||
return result
|
||||
return {"success": True, "updated_fields": ["model_name"]}
|
||||
|
||||
monkeypatch.setattr("py.services.agent.post_processor.PostProcessor", _Processor)
|
||||
return calls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orchestration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHydrateFromSource:
|
||||
@pytest.mark.asyncio
|
||||
async def test_applies_the_site_card_without_an_llm(self, monkeypatch):
|
||||
source = _FakeSource(
|
||||
context=ModelCardContext(
|
||||
model_name="Krea-2-LORA",
|
||||
version_name="c1-st1000",
|
||||
description="权重0.5-1.2。",
|
||||
official_tags=["photography"],
|
||||
),
|
||||
readme="# Krea-2-LORA",
|
||||
)
|
||||
calls = _wire(monkeypatch, source)
|
||||
|
||||
updated = await hydrate_from_source("/models/lora.safetensors", ref=REF)
|
||||
|
||||
assert updated == ["model_name"]
|
||||
assert len(calls) == 1
|
||||
call = calls[0]
|
||||
# No provider is consulted: everything applied is what the site published.
|
||||
assert call["llm_output"] == {}
|
||||
assert call["skill_name"] == "enrich_hf_metadata"
|
||||
assert call["readme_content"] == "# Krea-2-LORA"
|
||||
assert call["source_context"].model_name == "Krea-2-LORA"
|
||||
assert call["metadata_source"] == "source:modelscope"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_matches_the_file_by_hash_and_basename(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
_wire(monkeypatch, source)
|
||||
|
||||
await hydrate_from_source("/models/sub/Krea-2-LORA_c1-st1000.safetensors", ref=REF)
|
||||
|
||||
assert source.context_kwargs == {
|
||||
"filename": "Krea-2-LORA_c1-st1000.safetensors",
|
||||
"sha256": "a" * 64,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_for_an_unknown_platform(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(monkeypatch, source)
|
||||
monkeypatch.setattr(hydration, "get_source", lambda _platform: None)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_for_a_link_only_source(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
source.supports_enrichment = False
|
||||
calls = _wire(monkeypatch, source)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_without_a_sidecar(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(monkeypatch, source, metadata={})
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_when_the_site_published_nothing(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(), readme="")
|
||||
calls = _wire(monkeypatch, source)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deferred_hash_still_matches_by_filename(self, monkeypatch):
|
||||
"""Checkpoints and other large files are stored with
|
||||
``hash_status="pending"`` and an empty ``sha256`` (see
|
||||
``CheckpointScanner._create_default_metadata``), so hydration has to
|
||||
work from the filename alone."""
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(
|
||||
monkeypatch,
|
||||
source,
|
||||
metadata={**SIDECAR, "sha256": "", "hash_status": "pending"},
|
||||
)
|
||||
|
||||
await hydrate_from_source("/models/big_checkpoint.safetensors", ref=REF)
|
||||
|
||||
assert calls[0]["source_context"].model_name == "X"
|
||||
assert source.context_kwargs == {
|
||||
"filename": "big_checkpoint.safetensors",
|
||||
"sha256": "",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_when_the_model_is_not_linked(self, monkeypatch):
|
||||
"""A file that merely shares a name must not get another model's card."""
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(
|
||||
monkeypatch, source, metadata={"sha256": "a" * 64, "base_model": "Unknown"}
|
||||
)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_when_linked_to_another_repository(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(
|
||||
monkeypatch,
|
||||
source,
|
||||
metadata={
|
||||
**SIDECAR,
|
||||
"source_url": "https://modelscope.cn/models/user/other",
|
||||
},
|
||||
)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_early_when_linked_to_another_platform(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
calls = _wire(
|
||||
monkeypatch,
|
||||
source,
|
||||
metadata={
|
||||
"sha256": "a" * 64,
|
||||
"source_platform": "huggingface",
|
||||
"source_url": "https://huggingface.co/user/repo",
|
||||
},
|
||||
)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_readme_alone_is_enough_to_run(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(), readme="# hi")
|
||||
calls = _wire(monkeypatch, source)
|
||||
|
||||
await hydrate_from_source("/models/lora.safetensors", ref=REF)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert calls[0]["readme_content"] == "# hi"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_site_never_breaks_the_download(self, monkeypatch):
|
||||
source = _FakeSource(fail=True)
|
||||
calls = _wire(monkeypatch, source)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
assert calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_post_processor_never_breaks_the_download(
|
||||
self, monkeypatch
|
||||
):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
_wire(monkeypatch, source, result={"success": False, "errors": ["boom"]})
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_updates_are_reported_for_logging(self, monkeypatch):
|
||||
source = _FakeSource(context=ModelCardContext(model_name="X"))
|
||||
_wire(
|
||||
monkeypatch,
|
||||
source,
|
||||
result={"success": True, "updated_fields": ["tags", "civitai"]},
|
||||
)
|
||||
|
||||
assert await hydrate_from_source("/models/lora.safetensors", ref=REF) == [
|
||||
"tags",
|
||||
"civitai",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-repository memo
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSharedSourceCache:
|
||||
def test_same_repository_reuses_one_memo(self):
|
||||
assert shared_source_cache("modelscope", "u/r") is shared_source_cache(
|
||||
"modelscope", "u/r"
|
||||
)
|
||||
|
||||
def test_different_repositories_get_different_memos(self):
|
||||
assert shared_source_cache("modelscope", "u/r") is not shared_source_cache(
|
||||
"modelscope", "u/other"
|
||||
)
|
||||
|
||||
def test_entry_expires(self, monkeypatch):
|
||||
clock = {"now": 1000.0}
|
||||
monkeypatch.setattr(hydration.time, "monotonic", lambda: clock["now"])
|
||||
|
||||
first = shared_source_cache("modelscope", "u/r")
|
||||
clock["now"] += hydration.SHARED_CACHE_TTL + 1
|
||||
|
||||
assert shared_source_cache("modelscope", "u/r") is not first
|
||||
|
||||
def test_cache_is_bounded(self):
|
||||
for index in range(SHARED_CACHE_MAX_ENTRIES + 5):
|
||||
shared_source_cache("modelscope", f"u/r{index}")
|
||||
|
||||
assert len(hydration._shared_caches) == SHARED_CACHE_MAX_ENTRIES
|
||||
|
||||
|
||||
class TestLoadModelCard:
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_read_is_memoised(self):
|
||||
source = _FakeSource(readme="# hi")
|
||||
cache = ModelSourceCache()
|
||||
|
||||
assert await load_model_card(source, "u/r", cache) == "# hi"
|
||||
assert await load_model_card(source, "u/r", cache) == "# hi"
|
||||
assert source.readme_calls == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_read_is_retried(self):
|
||||
"""A transient failure must not be cached as "this repo has no card"."""
|
||||
source = _FakeSource(readme="")
|
||||
cache = ModelSourceCache()
|
||||
|
||||
await load_model_card(source, "u/r", cache)
|
||||
await load_model_card(source, "u/r", cache)
|
||||
|
||||
assert source.readme_calls == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_works_without_a_cache(self):
|
||||
source = _FakeSource(readme="# hi")
|
||||
|
||||
assert await load_model_card(source, "u/r") == "# hi"
|
||||
assert source.readme_calls == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Base-model resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolveSiteBaseModel:
|
||||
@pytest.mark.asyncio
|
||||
async def test_maps_the_sites_own_vocabulary(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"py.metadata_ops.list_base_models",
|
||||
_async(["Krea 2", "Flux.1 D"]),
|
||||
)
|
||||
|
||||
context = ModelCardContext(
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
base_model_aliases=["KREA_2_TURBO", "krea/Krea-2-Turbo"],
|
||||
)
|
||||
|
||||
assert await resolve_site_base_model(context) == "Krea 2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_hint_defers_instead_of_guessing(self, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"py.metadata_ops.list_base_models", _async(["Flux.1 D"])
|
||||
)
|
||||
|
||||
context = ModelCardContext(base_model="something/else")
|
||||
|
||||
assert await resolve_site_base_model(context) == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_hints_needs_no_vocabulary_lookup(self, monkeypatch):
|
||||
async def _boom(*_args, **_kwargs): # pragma: no cover - must not run
|
||||
raise AssertionError("list_base_models should not be called")
|
||||
|
||||
monkeypatch.setattr("py.metadata_ops.list_base_models", _boom)
|
||||
|
||||
assert await resolve_site_base_model(ModelCardContext()) == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_vocabulary_failure_is_not_fatal(self, monkeypatch):
|
||||
async def _boom(*_args, **_kwargs):
|
||||
raise RuntimeError("civitai down")
|
||||
|
||||
monkeypatch.setattr("py.metadata_ops.list_base_models", _boom)
|
||||
|
||||
assert await resolve_site_base_model(ModelCardContext(base_model="x")) == ""
|
||||
Reference in New Issue
Block a user