mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -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:
@@ -0,0 +1,235 @@
|
||||
"""Deterministic metadata hydration for freshly downloaded source models.
|
||||
|
||||
A CivitAI download writes a fully-populated metadata sidecar as part of the
|
||||
download itself: the name, the description, the tags, the trigger words and
|
||||
the example images all arrive with the file. A download from an external
|
||||
model source (ModelScope, Hugging Face) has the same information behind a
|
||||
public API, but historically landed as a bare filename plus a source URL that
|
||||
the user had to enrich by hand ("Enrich Metadata with AI").
|
||||
|
||||
This module closes that gap without involving an LLM. It fetches the linked
|
||||
site's model card, hands it to the same :class:`~py.services.agent.post_processor.PostProcessor`
|
||||
the AI skill uses, and writes the result. Everything it applies is data the
|
||||
site published, so it is safe to run automatically on every download and to
|
||||
treat as a fallback for the gaps the LLM would otherwise fill.
|
||||
|
||||
Nothing here may break a download: every failure is logged and normalised to
|
||||
"the site had nothing to contribute".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from .base import ModelCardContext, ModelSourceCache
|
||||
from .registry import get_source, resolve_source_ref
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from .base import ModelSource, SourceRef
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: How long a fetched repository payload stays usable. A download batch walks
|
||||
#: a repository's files one HTTP request at a time, and the README plus the
|
||||
#: detail payload describe the *repository*, not the file, so re-fetching them
|
||||
#: per file would be pure waste. They expire so an edited model card is still
|
||||
#: picked up by the next batch.
|
||||
SHARED_CACHE_TTL = 300.0
|
||||
|
||||
#: Upper bound on memoised repositories; a long-running server must not grow
|
||||
#: without limit.
|
||||
SHARED_CACHE_MAX_ENTRIES = 32
|
||||
|
||||
#: ``"<platform>:<source_id>"`` → ``(expiry, memo)``.
|
||||
_shared_caches: dict[str, tuple[float, ModelSourceCache]] = {}
|
||||
|
||||
|
||||
def shared_source_cache(platform: str, source_id: str) -> ModelSourceCache:
|
||||
"""Return a short-lived per-repository memo for download-time hydration."""
|
||||
|
||||
now = time.monotonic()
|
||||
key = f"{platform}:{source_id}"
|
||||
entry = _shared_caches.get(key)
|
||||
if entry is not None and entry[0] > now:
|
||||
return entry[1]
|
||||
|
||||
for expired in [k for k, (expiry, _) in _shared_caches.items() if expiry <= now]:
|
||||
_shared_caches.pop(expired, None)
|
||||
if len(_shared_caches) >= SHARED_CACHE_MAX_ENTRIES:
|
||||
oldest = min(_shared_caches, key=lambda k: _shared_caches[k][0])
|
||||
_shared_caches.pop(oldest, None)
|
||||
|
||||
cache = ModelSourceCache()
|
||||
_shared_caches[key] = (now + SHARED_CACHE_TTL, cache)
|
||||
return cache
|
||||
|
||||
|
||||
def reset_shared_caches() -> None:
|
||||
"""Drop every memoised repository — used by tests."""
|
||||
|
||||
_shared_caches.clear()
|
||||
|
||||
|
||||
async def load_model_card(
|
||||
source: "ModelSource",
|
||||
source_id: str,
|
||||
cache: Optional[ModelSourceCache] = None,
|
||||
) -> str:
|
||||
"""Return *source_id*'s README, reusing *cache* when one is supplied.
|
||||
|
||||
Only successful reads are memoised, leaving a transient failure to be
|
||||
retried for the next file of the same repository.
|
||||
"""
|
||||
|
||||
key = f"{source.platform}:{source_id}"
|
||||
if cache is not None:
|
||||
cached = cache.readmes.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
readme = await source.fetch_model_card(source_id)
|
||||
if cache is not None and readme:
|
||||
cache.readmes[key] = readme
|
||||
return readme or ""
|
||||
|
||||
|
||||
async def resolve_site_base_model(context: ModelCardContext) -> str:
|
||||
"""Resolve the site's base-model hints to a canonical name, or ``""``.
|
||||
|
||||
Sites name base models in their own vocabulary (ModelScope publishes both
|
||||
``krea/Krea-2-Turbo`` and the ``KREA_2_TURBO`` enum). The resolver is
|
||||
strict and only ever returns a name the canonical vocabulary already
|
||||
contains, so an uncertain hint yields ``""`` rather than a plausible-looking
|
||||
wrong value.
|
||||
"""
|
||||
|
||||
hints = [*context.base_model_aliases, context.base_model]
|
||||
if not any(hints):
|
||||
return ""
|
||||
|
||||
# Imported lazily: pulling in the agent package at module scope would make
|
||||
# the model-source package import itself while it is still initialising.
|
||||
try:
|
||||
from ...metadata_ops import list_base_models
|
||||
from ..agent.base_model_resolver import resolve_base_model
|
||||
|
||||
known_names = await list_base_models()
|
||||
except Exception as exc:
|
||||
logger.warning("Could not resolve a site base model: %s", exc)
|
||||
return ""
|
||||
return resolve_base_model(hints, known_names)
|
||||
|
||||
|
||||
async def hydrate_from_source(
|
||||
file_path: str,
|
||||
*,
|
||||
ref: "SourceRef",
|
||||
cache: Optional[ModelSourceCache] = None,
|
||||
) -> list[str]:
|
||||
"""Apply the linked site's published metadata to a downloaded model.
|
||||
|
||||
This is the deterministic counterpart of the ``enrich_hf_metadata`` skill:
|
||||
it produces the same populated model card a CivitAI download produces,
|
||||
without an LLM and without user action.
|
||||
|
||||
Args:
|
||||
file_path: The just-downloaded model file, whose sidecar already
|
||||
carries the SHA256 used to match the right file in a collection
|
||||
repository.
|
||||
ref: The source the file came from.
|
||||
cache: Optional per-call memo; defaults to a short-lived shared one so
|
||||
a batch over one repository fetches its card only once.
|
||||
|
||||
Returns:
|
||||
The names of the metadata fields that changed. Never raises — a site
|
||||
that is down, or an API that changed shape, must not fail a download.
|
||||
"""
|
||||
|
||||
try:
|
||||
source = get_source(ref.platform)
|
||||
if source is None or not source.supports_enrichment:
|
||||
return []
|
||||
|
||||
from ...metadata_ops import read_metadata
|
||||
|
||||
metadata = await read_metadata(file_path)
|
||||
if not metadata:
|
||||
logger.debug("No metadata to hydrate for %s", file_path)
|
||||
return []
|
||||
|
||||
# Only a model that is actually linked to this repository may be
|
||||
# updated. The download path writes those fields just before calling
|
||||
# us; a file that merely shares a name with the requested one must not
|
||||
# be given another model's card.
|
||||
linked = resolve_source_ref(metadata)
|
||||
if linked is None or (linked.platform, linked.source_id) != (
|
||||
ref.platform,
|
||||
ref.source_id,
|
||||
):
|
||||
logger.debug(
|
||||
"Not hydrating %s: linked to %s, not %s",
|
||||
file_path, linked.url if linked else "no model source", ref.url,
|
||||
)
|
||||
return []
|
||||
|
||||
memo = cache if cache is not None else shared_source_cache(
|
||||
ref.platform, ref.source_id
|
||||
)
|
||||
readme = await load_model_card(source, ref.source_id, memo)
|
||||
context = await source.fetch_model_card_context(
|
||||
ref.source_id,
|
||||
os.path.basename(file_path),
|
||||
sha256=(metadata.get("sha256") or "").strip(),
|
||||
cache=memo,
|
||||
)
|
||||
if context.is_empty() and not readme:
|
||||
logger.debug(
|
||||
"No published metadata for %s on %s", ref.source_id, ref.platform
|
||||
)
|
||||
return []
|
||||
|
||||
resolved_base_model = await resolve_site_base_model(context)
|
||||
|
||||
from ..agent.post_processor import PostProcessor
|
||||
|
||||
result = await PostProcessor().process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path=file_path,
|
||||
llm_output={},
|
||||
metadata=metadata,
|
||||
readme_content=readme,
|
||||
source_context=context,
|
||||
resolved_base_model=resolved_base_model,
|
||||
metadata_source=f"source:{ref.platform}",
|
||||
)
|
||||
if not result.get("success", True):
|
||||
logger.debug(
|
||||
"Hydration reported failure for %s: %s",
|
||||
file_path, result.get("errors"),
|
||||
)
|
||||
return []
|
||||
|
||||
updated = list(result.get("updated_fields") or [])
|
||||
logger.info(
|
||||
"Hydrated %s from %s (%s): %s",
|
||||
file_path, source.label or ref.platform, ref.source_id,
|
||||
", ".join(updated) or "nothing to change",
|
||||
)
|
||||
return updated
|
||||
except Exception as exc: # pragma: no cover - defensive by design
|
||||
logger.warning("Source hydration failed for %s: %s", file_path, exc)
|
||||
return []
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SHARED_CACHE_MAX_ENTRIES",
|
||||
"SHARED_CACHE_TTL",
|
||||
"hydrate_from_source",
|
||||
"load_model_card",
|
||||
"reset_shared_caches",
|
||||
"resolve_site_base_model",
|
||||
"shared_source_cache",
|
||||
]
|
||||
Reference in New Issue
Block a user