alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-21 03:01:27 -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/services/test_model_sources.py
+185 -2
View File
@@ -59,6 +59,17 @@ class TestDetectSource:
"modelscope",
"jj3550945163/Krea-2-LORA",
),
# modelscope.ai is a separate catalogue with its own platform id.
(
"https://www.modelscope.ai/models/referall13/EM1",
"modelscope-ai",
"referall13/EM1",
),
(
"https://modelscope.ai/models/ErLubu/krea2_style_260911_02/summary",
"modelscope-ai",
"ErLubu/krea2_style_260911_02",
),
(
"https://tensor.art/models/827823520299086029/Vivid-Impressions-Storybook-Sstyle-V1.0",
"tensorart",
@@ -97,6 +108,21 @@ class TestDetectSource:
== "https://tensor.art/models/123"
)
def test_modelscope_com_is_an_alias_of_the_mainland_site(self):
"""``.com`` 301-redirects to ``.cn``, so it is not a third catalogue."""
ref = detect_source("https://www.modelscope.com/models/u/r")
assert ref.platform == "modelscope"
assert ref.url == "https://modelscope.cn/models/u/r"
def test_the_two_modelscope_catalogues_do_not_cross_match(self):
"""A host must never be accepted by the other deployment's patterns."""
mainland = get_source("modelscope")
international = get_source("modelscope-ai")
assert mainland.parse("https://www.modelscope.ai/models/u/r") is None
assert international.parse("https://modelscope.cn/models/u/r") is None
assert international.parse("https://www.modelscope.com/models/u/r") is None
class TestStrictParsing:
@pytest.mark.parametrize(
@@ -106,6 +132,8 @@ class TestStrictParsing:
"https://huggingface.co/user/repo/",
"https://modelscope.cn/models/user/repo",
"https://modelscope.cn/models/user/repo/summary",
"https://www.modelscope.ai/models/user/repo",
"https://www.modelscope.ai/models/user/repo/files",
"https://tensor.art/models/827823520299086029",
"https://tensor.art/models/827823520299086029/Vivid-Impressions",
],
@@ -145,6 +173,16 @@ class TestCapabilities:
assert source.default_revision == "master"
assert source.default_subdir == "modelscope"
def test_modelscope_intl_is_the_same_site_on_another_catalogue(self):
source = get_source("modelscope-ai")
assert source.supports_enrichment is True
assert source.supports_download is True
assert source.default_revision == "master"
# A distinct directory: the same owner/name can exist on both
# deployments with different content.
assert source.default_subdir == "modelscope-ai"
assert source.base_url == "https://www.modelscope.ai"
def test_tensorart_is_link_only(self):
source = get_source("tensorart")
assert source.supports_enrichment is False
@@ -152,11 +190,17 @@ class TestCapabilities:
def test_registry_lists_every_source(self):
platforms = {s.platform for s in list_sources()}
assert platforms == {"huggingface", "modelscope", "tensorart"}
assert platforms == {
"huggingface",
"modelscope",
"modelscope-ai",
"tensorart",
}
def test_labels_are_brand_names(self):
assert source_label("huggingface") == "Hugging Face"
assert source_label("modelscope") == "ModelScope"
assert source_label("modelscope-ai") == "ModelScope (International)"
assert source_label("tensorart") == "TensorArt"
assert source_label("unknown", "fallback") == "fallback"
@@ -311,6 +355,44 @@ class TestFetchModelCard:
in calls
)
@pytest.mark.asyncio
async def test_modelscope_intl_fetches_from_its_own_catalogue(self, monkeypatch):
"""The mainland site 404s for a `.ai`-only repository, so every fetch
has to stay on the host the URL came from."""
calls: list[str] = []
async def fake_fetch_text(url: str, **_kwargs) -> str:
calls.append(url)
return "# card"
json_calls: list[str] = []
async def fake_fetch_json(url: str, **_kwargs):
json_calls.append(url)
return 200, {"Data": {"Name": "EM1", "MuseInfo": {"versions": []}}}
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_text", fake_fetch_text
)
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
source = get_source("modelscope-ai")
await source.fetch_model_card("referall13/EM1")
await source.fetch_model_card_context("referall13/EM1")
await source.list_files("referall13/EM1")
assert calls == [
"https://www.modelscope.ai/models/referall13/EM1/resolve/master/README.md"
]
assert json_calls == [
"https://www.modelscope.ai/api/v1/models/referall13/EM1",
"https://www.modelscope.ai/api/v1/models/referall13/EM1/repo/files"
"?Revision=master",
]
assert not any("modelscope.cn" in url for url in calls + json_calls)
@pytest.mark.asyncio
async def test_tensorart_never_fetches(self):
# TensorArt enrichment is disabled: the provider must not issue any
@@ -335,6 +417,12 @@ class TestAssetBaseUrl:
== "https://modelscope.cn/models/u/r/resolve/master"
)
def test_modelscope_intl_uses_master_revision(self):
assert (
get_source("modelscope-ai").asset_base_url("u/r")
== "https://www.modelscope.ai/models/u/r/resolve/master"
)
# ---------------------------------------------------------------------------
# Model card context (site extras kept outside the README)
@@ -354,6 +442,9 @@ def _modelscope_detail_payload() -> dict:
"Data": {
"Name": "Krea-2-LORA",
"ChineseName": "krea脸模",
"AigcType": "LoRA",
"License": "Apache License 2.0",
"Tags": ["LoRA", "text-to-image", "portrait"],
"Description": "权重0.5-1.2。配合《风格滤镜》lora一起使用。",
"BaseModel": ["krea/Krea-2-Turbo"],
"License": "Apache License 2.0",
@@ -422,6 +513,81 @@ class TestFetchModelCardContext:
# OfficialTag values only, de-duplicated, order preserved.
assert context.official_tags == ["photography", "woman"]
@pytest.mark.asyncio
async def test_modelscope_reads_site_identity_fields(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 200, _modelscope_detail_payload()
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
context = await ModelScopeSource().fetch_model_card_context(
"u/r", "Krea-2-LORA_c1-st1000.safetensors"
)
assert context.model_name == "Krea-2-LORA"
assert context.model_name_localized == "krea脸模"
assert context.license == "Apache License 2.0"
assert context.model_type == "LoRA"
# The version label is taken from the file that was matched, not from
# whichever version happens to come first in the payload.
assert context.version_name == "c1-st1000"
@pytest.mark.asyncio
async def test_modelscope_version_label_is_empty_for_an_unknown_file(
self, monkeypatch
):
async def fake_fetch_json(url, **_kwargs):
return 200, _modelscope_detail_payload()
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
context = await ModelScopeSource().fetch_model_card_context(
"u/r", "other.safetensors"
)
assert context.version_name == ""
# The repository-wide fields are still published.
assert context.model_name == "Krea-2-LORA"
@pytest.mark.asyncio
async def test_modelscope_falls_back_to_plain_tags(self, monkeypatch):
"""An empty ``OfficialTags`` must not mean "no tags at all".
The plain ``Tags`` list mixes genuine content tags with library and
task categories; the latter are dropped so the card is not tagged
"lora" / "text-to-image".
"""
payload = _modelscope_detail_payload()
payload["Data"]["OfficialTags"] = None
async def fake_fetch_json(url, **_kwargs):
return 200, payload
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
context = await ModelScopeSource().fetch_model_card_context("u/r")
assert context.official_tags == ["portrait"]
@pytest.mark.asyncio
async def test_modelscope_curated_tags_win_over_plain_tags(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
return 200, _modelscope_detail_payload()
monkeypatch.setattr(
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
)
context = await ModelScopeSource().fetch_model_card_context("u/r")
assert "portrait" not in context.official_tags
@pytest.mark.asyncio
async def test_modelscope_matches_example_images_by_filename(self, monkeypatch):
async def fake_fetch_json(url, **_kwargs):
@@ -655,6 +821,22 @@ class TestDownloadUrls:
"https://modelscope.cn/models/u/r/resolve/master/sub/f.safetensors"
)
def test_modelscope_intl_builds_every_url_on_its_own_host(self):
"""The two deployments serve different catalogues, so a URL built for
one must never point at the other."""
source = get_source("modelscope-ai")
assert source.canonical_url("u/r") == "https://www.modelscope.ai/models/u/r"
assert source.file_download_url("u/r", "sub/f.safetensors") == (
"https://www.modelscope.ai/models/u/r/resolve/master/sub/f.safetensors"
)
assert source.asset_base_url("u/r") == (
"https://www.modelscope.ai/models/u/r/resolve/master"
)
assert source.page_url_for_file("u/r", "sub/f.safetensors") == (
"https://www.modelscope.ai/models/u/r/file/view/master/sub/f.safetensors"
)
def test_explicit_revision_wins(self):
assert ModelScopeSource().file_download_url("u/r", "f.bin", "v1") == (
"https://modelscope.cn/models/u/r/resolve/v1/f.bin"
@@ -701,12 +883,13 @@ class TestSourceIdValidation:
class TestDownloadSourceRegistry:
def test_downloadable_sources_excludes_link_only_sites(self):
platforms = {source.platform for source in downloadable_sources()}
assert platforms == {"huggingface", "modelscope"}
assert platforms == {"huggingface", "modelscope", "modelscope-ai"}
def test_get_download_source_rejects_link_only_platform(self):
assert get_download_source("tensorart") is None
assert get_download_source("nope") is None
assert get_download_source("modelscope").platform == "modelscope"
assert get_download_source("modelscope-ai").platform == "modelscope-ai"
assert get_download_source("huggingface").platform == "huggingface"
tests/services/test_modelscope_hydration.py
+397
View File
@@ -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")) == ""
tests/services/test_post_processor.py
+115
View File
@@ -1121,3 +1121,118 @@ pip install modelscope
)
assert "modelDescription" not in mock_apply.call_args[0][1]
# ======================================================================
# Site identity and provenance fields
# ======================================================================
class TestSiteIdentityFields:
"""The fields that make a source download look like a CivitAI one."""
METADATA = {
"from_civitai": False,
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/user/repo",
"file_name": "Krea-2-LORA_c1-st1000",
"model_name": "Krea-2-LORA_c1-st1000",
"base_model": "Unknown",
}
@staticmethod
def _run(processor, *, metadata, context, llm_output=None, **kwargs):
async def _call():
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm_output if llm_output is not None else {},
metadata=metadata,
source_context=context,
**kwargs,
)
return mock_apply.call_args[0][1]
return _call()
@pytest.mark.asyncio
async def test_model_name_is_taken_from_the_site(self, processor):
applied = await self._run(
processor,
metadata=dict(self.METADATA),
context=ModelCardContext(model_name="Krea-2-LORA"),
)
assert applied["model_name"] == "Krea-2-LORA"
@pytest.mark.asyncio
async def test_model_name_is_also_written_when_absent(self, processor):
metadata = {**self.METADATA, "model_name": ""}
applied = await self._run(
processor,
metadata=metadata,
context=ModelCardContext(model_name="Krea-2-LORA"),
)
assert applied["model_name"] == "Krea-2-LORA"
@pytest.mark.asyncio
async def test_renamed_model_keeps_the_users_name(self, processor):
metadata = {**self.METADATA, "model_name": "my own name"}
applied = await self._run(
processor,
metadata=metadata,
context=ModelCardContext(model_name="Krea-2-LORA"),
)
assert "model_name" not in applied
@pytest.mark.asyncio
async def test_version_label_becomes_the_civitai_name(self, processor):
applied = await self._run(
processor,
metadata=dict(self.METADATA),
context=ModelCardContext(
model_name="Krea-2-LORA",
version_name="c1-st1000",
description="权重0.5-1.2。",
),
)
assert applied["civitai"]["name"] == "c1-st1000"
# Every civitai branch contributes to one dict, so an earlier branch
# must survive a later one.
assert applied["civitai"]["description"] == "权重0.5-1.2。"
@pytest.mark.asyncio
async def test_llm_enriched_at_is_stamped_only_when_the_llm_answered(
self, processor
):
applied = await self._run(
processor,
metadata=dict(self.METADATA),
context=ModelCardContext(model_name="Krea-2-LORA"),
)
assert applied["metadata_source"] == "agent:enrich_hf_metadata"
assert "llm_enriched_at" not in applied
@pytest.mark.asyncio
async def test_llm_answer_stamps_llm_enriched_at(self, processor):
applied = await self._run(
processor,
metadata=dict(self.METADATA),
context=ModelCardContext(model_name="Krea-2-LORA"),
llm_output={"base_model": "", "confidence": "high"},
)
assert "llm_enriched_at" in applied
@pytest.mark.asyncio
async def test_metadata_source_can_be_overridden(self, processor):
applied = await self._run(
processor,
metadata=dict(self.METADATA),
context=ModelCardContext(model_name="Krea-2-LORA"),
metadata_source="source:modelscope",
)
assert applied["metadata_source"] == "source:modelscope"