alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-20 18:51:26 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(download): fill model metadata from the source API on download

Browse Source 
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
2026-09-17 07:47:07 +08:00
parent 1b1a8d63db
commit d572292142
34 changed files with 2385 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/routes/test_model_source_handlers.py
+521 -6
View File
@@ -307,7 +307,12 @@ async def test_get_model_sources_lists_capabilities():
sources = _json_payload(response)
by_platform = {s["platform"]: s for s in sources}
assert set(by_platform) == {"huggingface", "modelscope", "tensorart"}
assert set(by_platform) == {
"huggingface",
"modelscope",
"modelscope-ai",
"tensorart",
}
assert by_platform["huggingface"]["supports_enrichment"] is True
assert by_platform["modelscope"]["supports_enrichment"] is True
# TensorArt is link-only: no accessible model card for the backend.
@@ -315,6 +320,12 @@ async def test_get_model_sources_lists_capabilities():
assert by_platform["modelscope"]["supports_download"] is True
assert by_platform["modelscope"]["default_revision"] == "master"
assert by_platform["tensorart"]["supports_download"] is False
# The international deployment is advertised with its own example URL, so
# the Link dialog names the host a user actually has open.
assert by_platform["modelscope-ai"]["supports_download"] is True
assert by_platform["modelscope-ai"]["example_url"].startswith(
"https://www.modelscope.ai/"
)
assert all(s["example_url"] for s in sources)
@@ -492,6 +503,44 @@ async def test_download_model_source_modelscope_default_paths(tmp_path, monkeypa
)
@pytest.mark.asyncio
async def test_download_model_source_modelscope_intl_uses_its_own_host(
tmp_path, monkeypatch
):
"""`.ai` is a separate catalogue, so the download must not go to `.cn`."""
captured = _stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope-ai",
"repo": "referall13/EM1",
"filename": "EM1_c1-st1000.safetensors",
"model_root": str(tmp_path),
"use_default_paths": True,
}
)
)
assert response.status == 200
assert captured["url"] == (
"https://www.modelscope.ai/models/referall13/EM1/resolve/master/"
"EM1_c1-st1000.safetensors"
)
# Its own default directory, so the same owner/name on both deployments
# cannot overwrite each other.
assert captured["save_path"] == str(
tmp_path / "modelscope-ai" / "referall13" / "EM1" / "EM1_c1-st1000.safetensors"
)
ref = saved.await_args.args[1]
assert ref.platform == "modelscope-ai"
assert ref.source_id == "referall13/EM1"
assert ref.url == "https://www.modelscope.ai/models/referall13/EM1"
@pytest.mark.asyncio
async def test_download_model_source_defaults_to_huggingface(tmp_path, monkeypatch):
"""The legacy /api/lm/download-hf-model payload has no `platform` key."""
@@ -609,18 +658,19 @@ async def test_save_source_metadata_writes_platform_fields(
base_model="SDXL 1.0",
preview_url="",
)
monkeypatch.setattr(
model_source_handlers.MetadataManager,
"create_default_metadata",
AsyncMock(return_value=metadata),
scanner = SimpleNamespace(
# A real scanner owns metadata creation (see the lazy-hash test below).
_create_default_metadata=AsyncMock(return_value=metadata),
add_model_to_cache=AsyncMock(),
)
scanner = SimpleNamespace(add_model_to_cache=AsyncMock())
monkeypatch.setattr(
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=scanner)
)
monkeypatch.setattr(
model_source_handlers, "_infer_model_type", lambda _root: (LoraMetadata, "get_lora_scanner")
)
hydrate = AsyncMock()
monkeypatch.setattr(model_source_handlers, "hydrate_from_source", hydrate)
ref = SourceRef(platform=platform, source_id="u/r", url=url)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
@@ -630,6 +680,471 @@ async def test_save_source_metadata_writes_platform_fields(
assert saved["source_url"] == url
assert bool(saved.get("hf_url", "")) is expect_hf_alias
assert scanner._create_default_metadata.await_args.args == (str(model_path),)
cached = scanner.add_model_to_cache.await_args.args[0]
assert cached["source_platform"] == platform
assert cached["source_url"] == url
# The site's own API is consulted last, so the scanner-cache refresh it
# performs lands on the entry created above.
assert hydrate.await_args.args == (str(model_path),)
assert hydrate.await_args.kwargs["ref"] == ref
@pytest.mark.asyncio
async def test_checkpoint_download_defers_the_hash(tmp_path, monkeypatch):
"""A multi-GB checkpoint must not be hashed inside the download request.
``CheckpointScanner`` records ``hash_status="pending"`` and lets the hash be
computed on demand; going through the generic
``MetadataManager.create_default_metadata`` would read the whole file before
the download response could return, which is exactly the pause this code
path is supposed to avoid.
"""
from py.services.checkpoint_scanner import CheckpointScanner
from py.utils.models import CheckpointMetadata
model_path = tmp_path / "big_checkpoint.safetensors"
model_path.write_bytes(b"stub")
real_scanner = CheckpointScanner()
scanner = SimpleNamespace(
_create_default_metadata=real_scanner._create_default_metadata,
add_model_to_cache=AsyncMock(),
)
monkeypatch.setattr(
ServiceRegistry, "get_checkpoint_scanner", AsyncMock(return_value=scanner)
)
monkeypatch.setattr(
model_source_handlers,
"_infer_model_type",
lambda _root: (CheckpointMetadata, "get_checkpoint_scanner"),
)
generic = AsyncMock(
# Stands in for the eager helper: if the handler reaches for it, the
# sidecar ends up hashed and the assertions below say so plainly.
return_value=LoraMetadata(
file_name="big_checkpoint",
model_name="big_checkpoint",
file_path=str(model_path),
size=4,
modified=1.0,
sha256="d" * 64,
base_model="Unknown",
preview_url="",
)
)
monkeypatch.setattr(
model_source_handlers.MetadataManager, "create_default_metadata", generic
)
monkeypatch.setattr(model_source_handlers, "hydrate_from_source", AsyncMock())
ref = SourceRef(
platform="huggingface",
source_id="u/r",
url="https://huggingface.co/u/r",
)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["sha256"] == ""
assert saved["hash_status"] == "pending"
assert saved["from_civitai"] is False
# The download link is still recorded on top of the deferred hash.
assert saved["source_platform"] == "huggingface"
assert saved["source_url"] == "https://huggingface.co/u/r"
# The scanner cache must carry the pending state too, or the cache fill
# would compute the hash after all.
cached = scanner.add_model_to_cache.await_args.args[0]
assert cached["hash_status"] == "pending"
assert cached["sha256"] == ""
generic.assert_not_awaited()
# ---------------------------------------------------------------------------
# Post-transfer phase reporting
# ---------------------------------------------------------------------------
def _stub_hydration_pipeline(tmp_path, monkeypatch):
"""Wire `_save_source_metadata`'s collaborators and record call order."""
model_path = tmp_path / "downloaded.safetensors"
model_path.write_bytes(b"x" * 32)
metadata = LoraMetadata(
file_name="downloaded",
model_name="Downloaded",
file_path=str(model_path),
size=32,
modified=1.0,
sha256="a" * 64,
base_model="SDXL 1.0",
preview_url="",
)
monkeypatch.setattr(
model_source_handlers.MetadataManager,
"create_default_metadata",
AsyncMock(return_value=metadata),
)
monkeypatch.setattr(
ServiceRegistry,
"get_lora_scanner",
AsyncMock(return_value=SimpleNamespace(add_model_to_cache=AsyncMock())),
)
monkeypatch.setattr(
model_source_handlers,
"_infer_model_type",
lambda _root: (LoraMetadata, "get_lora_scanner"),
)
events: list = []
async def fake_broadcast(download_id, data):
events.append(("broadcast", data["stage"], data, download_id))
async def fake_hydrate(*_args, **_kwargs):
events.append(("hydrate", None, None, None))
monkeypatch.setattr(
model_source_handlers.ws_manager, "broadcast_download_progress", fake_broadcast
)
monkeypatch.setattr(model_source_handlers, "hydrate_from_source", fake_hydrate)
return model_path, events
@pytest.mark.asyncio
async def test_save_source_metadata_reports_post_transfer_stages(tmp_path, monkeypatch):
"""The byte counter stops before indexing and the site fetch, so the UI has
to be told what is still running — otherwise the bar looks stuck."""
model_path, events = _stub_hydration_pipeline(tmp_path, monkeypatch)
ref = SourceRef(
platform="modelscope", source_id="u/r", url="https://modelscope.cn/models/u/r"
)
await model_source_handlers._save_source_metadata(
str(model_path), ref, str(tmp_path), download_id="dl-1"
)
# Each stage is announced *before* its work starts, so the label is never
# describing something that already finished.
assert [event[:2] for event in events] == [
("broadcast", "indexing"),
("broadcast", "source"),
("hydrate", None),
]
for kind, stage, data, download_id in events:
if kind != "broadcast":
continue
assert download_id == "dl-1"
assert data["status"] == "metadata"
assert data["progress"] == 100
assert data["platform"] == "modelscope"
@pytest.mark.asyncio
async def test_save_source_metadata_is_silent_without_a_watcher(tmp_path, monkeypatch):
"""No `download_id` means no UI is watching; nothing should be broadcast."""
model_path, events = _stub_hydration_pipeline(tmp_path, monkeypatch)
ref = SourceRef(
platform="modelscope", source_id="u/r", url="https://modelscope.cn/models/u/r"
)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
assert events == [("hydrate", None, None, None)]
@pytest.mark.asyncio
async def test_report_phase_never_breaks_a_download(monkeypatch):
"""Progress reporting is cosmetic; a dead socket must not fail the file."""
monkeypatch.setattr(
model_source_handlers.ws_manager,
"broadcast_download_progress",
AsyncMock(side_effect=RuntimeError("socket gone")),
)
await model_source_handlers._report_phase("dl-1", "source", "modelscope")
@pytest.mark.asyncio
async def test_download_passes_its_watch_id_into_metadata_work(tmp_path, monkeypatch):
"""The stages are only visible if the handler hands its id down."""
_stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "owner/name",
"filename": "model.safetensors",
"model_root": str(tmp_path),
"download_id": "dl-42",
}
)
)
assert saved.await_args.kwargs["download_id"] == "dl-42"
@pytest.mark.asyncio
async def test_skipped_download_still_reports_the_site_stage(tmp_path, monkeypatch):
"""An already-present file is hydrated too, so it needs the same signal."""
_stub_download_backend(monkeypatch)
hydrate = AsyncMock()
monkeypatch.setattr(model_source_handlers, "hydrate_from_source", hydrate)
broadcast = AsyncMock()
monkeypatch.setattr(
model_source_handlers.ws_manager, "broadcast_download_progress", broadcast
)
existing = tmp_path / "model.safetensors"
existing.write_bytes(b"x" * 32)
await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "owner/name",
"filename": "model.safetensors",
"model_root": str(tmp_path),
"download_id": "dl-7",
}
)
)
assert broadcast.await_args.args[1]["stage"] == "source"
@pytest.mark.asyncio
async def test_save_source_metadata_survives_a_hydration_failure(tmp_path, monkeypatch):
"""Metadata hydration must never be able to fail a completed download."""
model_path = tmp_path / "downloaded.safetensors"
model_path.write_bytes(b"x" * 32)
metadata = LoraMetadata(
file_name="downloaded",
model_name="Downloaded",
file_path=str(model_path),
size=32,
modified=1.0,
sha256="a" * 64,
base_model="SDXL 1.0",
preview_url="",
)
monkeypatch.setattr(
model_source_handlers.MetadataManager,
"create_default_metadata",
AsyncMock(return_value=metadata),
)
monkeypatch.setattr(
ServiceRegistry,
"get_lora_scanner",
AsyncMock(return_value=SimpleNamespace(add_model_to_cache=AsyncMock())),
)
monkeypatch.setattr(
model_source_handlers, "_infer_model_type", lambda _root: (LoraMetadata, "get_lora_scanner")
)
monkeypatch.setattr(
model_source_handlers,
"hydrate_from_source",
AsyncMock(side_effect=RuntimeError("site down")),
)
ref = SourceRef(
platform="modelscope", source_id="u/r", url="https://modelscope.cn/models/u/r"
)
await model_source_handlers._save_source_metadata(str(model_path), ref, str(tmp_path))
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
assert saved["source_platform"] == "modelscope"
@pytest.mark.asyncio
async def test_downloading_an_existing_file_still_hydrates(tmp_path, monkeypatch):
"""A pre-existing file may still be missing the site's metadata."""
_stub_download_backend(monkeypatch)
saved = AsyncMock()
monkeypatch.setattr(model_source_handlers, "_save_source_metadata", saved)
hydrate = AsyncMock()
monkeypatch.setattr(model_source_handlers, "hydrate_from_source", hydrate)
existing = tmp_path / "model.safetensors"
existing.write_bytes(b"x" * 32)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "owner/name",
"filename": "model.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
saved.assert_not_awaited()
assert hydrate.await_args.args == (str(existing),)
assert hydrate.await_args.kwargs["ref"].source_id == "owner/name"
# ---------------------------------------------------------------------------
# Download-time metadata hydration (end to end)
# ---------------------------------------------------------------------------
def _modelscope_card_payload() -> dict:
"""A trimmed ModelScope model-detail response for the hydration test."""
return {
"Code": 200,
"Data": {
"Name": "Krea-2-LORA",
"ChineseName": "krea脸模",
"AigcType": "LoRA",
"Description": "权重0.5-1.2。配合《风格滤镜》lora一起使用。",
"BaseModel": ["krea/Krea-2-Turbo"],
"OfficialTags": [{"Tag": "photography"}, {"Tag": "woman"}],
"ModelInfos": {
"safetensor": {
"files": [
{
"name": "Krea-2-LORA_c1-st1000.safetensors",
"sha256": "a" * 64,
}
]
}
},
"MuseInfo": {
"versions": [
{
"stats": {"fileList": ["Krea-2-LORA_c1-st1000.safetensors"]},
"modelVersion": {
"showName": "c1-st1000",
"triggerWords": '["kreaface","kreamodel"]',
},
"coverImages": [
{"url": "https://resources.modelscope.cn/cover-images/b.png"},
{"url": "https://resources.modelscope.cn/cover-images/c.png"},
],
}
]
},
},
}
@pytest.mark.asyncio
async def test_download_hydrates_the_card_from_the_site(tmp_path, monkeypatch):
"""A ModelScope download must land with a populated model card.
Only the network, the scanner and the file transfer are faked, so this
exercises the real handler, the real `ModelScopeSource` and the real
post-processor together. Breaking the wiring between them fails here even
when each half still passes its own unit tests.
"""
model_path = tmp_path / "Krea-2-LORA_c1-st1000.safetensors"
async def fake_download_file(**kwargs):
with open(kwargs["save_path"], "wb") as handle:
handle.write(b"stub")
return True, kwargs["save_path"]
class _Downloader:
download_file = staticmethod(fake_download_file)
class _Settings:
def get(self, key, default=None):
return default
async def fake_get_downloader():
return _Downloader()
monkeypatch.setattr(model_source_handlers, "get_downloader", fake_get_downloader)
monkeypatch.setattr(
model_source_handlers, "get_settings_manager", lambda: _Settings()
)
monkeypatch.setattr(
model_source_handlers,
"_infer_model_type",
lambda _root: (LoraMetadata, "get_lora_scanner"),
)
scanner = SimpleNamespace(
get_cached_data=AsyncMock(
return_value=SimpleNamespace(raw_data=[{"file_path": str(model_path)}])
),
add_model_to_cache=AsyncMock(),
update_single_model_cache=AsyncMock(),
)
monkeypatch.setattr(
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=scanner)
)
async def fake_fetch_text(url, **_kwargs):
return "# Krea-2-LORA\n\n权重0.5-1.2。"
async def fake_fetch_json(url, **_kwargs):
return 200, _modelscope_card_payload()
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_text", fake_fetch_text
)
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
monkeypatch.setattr(
"py.metadata_ops.list_base_models", AsyncMock(return_value=["Krea 2"])
)
monkeypatch.setattr(
"py.metadata_ops.download_preview",
AsyncMock(return_value=str(tmp_path / "preview.webp")),
)
response = await ModelSourceHandler().download_model_source(
FakeRequest(
json_data={
"platform": "modelscope",
"repo": "jj3550945163/Krea-2-LORA",
"filename": "Krea-2-LORA_c1-st1000.safetensors",
"model_root": str(tmp_path),
}
)
)
assert response.status == 200
saved = json.loads(open(_sidecar_path(model_path), encoding="utf-8").read())
# The download's own provenance is unchanged.
assert saved["source_platform"] == "modelscope"
assert saved["source_url"] == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA"
assert saved["from_civitai"] is False
# The site's published metadata, with no LLM involved.
assert saved["model_name"] == "Krea-2-LORA"
assert saved["base_model"] == "Krea 2"
assert saved["tags"] == ["photography", "woman"]
assert saved["civitai"]["name"] == "c1-st1000"
assert saved["civitai"]["trainedWords"] == ["kreaface", "kreamodel"]
assert saved["civitai"]["description"] == "权重0.5-1.2。配合《风格滤镜》lora一起使用。"
assert [img["url"] for img in saved["civitai"]["images"]] == [
"https://resources.modelscope.cn/cover-images/b.png",
"https://resources.modelscope.cn/cover-images/c.png",
]
assert saved["preview_url"] == str(tmp_path / "preview.webp")
assert saved["usage_tips"] == (
'{"strength_min": 0.5, "strength_max": 1.2, "strength_range": "0.5-1.2"}'
)
assert saved["metadata_source"] == "source:modelscope"
# No provider answered, so claiming an AI enrichment would be a lie.
assert "llm_enriched_at" not in saved
# The enriched card reaches the scanner cache, not just the file.
assert scanner.update_single_model_cache.await_count == 1
cached = scanner.update_single_model_cache.await_args.args[2]
assert cached["model_name"] == "Krea-2-LORA"