mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
Compare commits
8 Commits
db38ad80e6
...
942717f0b6
| Author | SHA1 | Date | |
|---|---|---|---|
| 942717f0b6 | |||
| 0f160e157f | |||
| e9e9ee20c6 | |||
| f0ee30fc68 | |||
| 51de85a6ca | |||
| 4064ea7d3a | |||
| 35b291ab19 | |||
| e711e643f1 |
+56
-4
@@ -78,19 +78,71 @@ TensorArt is link-only: `tensor.art` sits behind a Cloudflare managed challenge
|
||||
|
||||
**What it does**:
|
||||
1. Reads the model's `.metadata.json` to get the source (`source_platform` + `source_url`, or the legacy `hf_url`)
|
||||
2. Fetches the model card through the provider in `py/services/model_sources/`
|
||||
3. Sends the README + local metadata to the LLM for structured extraction
|
||||
2. Fetches the model card through the provider in `py/services/model_sources/` — the README via `fetch_model_card()`, plus any extras the site keeps outside it via `fetch_model_card_context()`
|
||||
3. Sends the README + site-provided extras + local metadata to the LLM for structured extraction
|
||||
4. Writes extracted fields to `.metadata.json`:
|
||||
- `base_model` — only if current value is empty
|
||||
- `trainedWords` — trigger words (LoRA only, if none exist)
|
||||
- `modelDescription` — concise summary (if none exists)
|
||||
- `modelDescription` — the site's author description (if any) followed by the README rendered as HTML
|
||||
- `tags` — merged with existing tags, deduplicated
|
||||
- `civitai.images` — example images
|
||||
- `metadata_source` — audit trail: `agent:enrich_hf_metadata`
|
||||
- `llm_enriched_at` — ISO timestamp
|
||||
5. Downloads and optimizes preview image (if LLM found one in the README)
|
||||
5. Downloads and optimizes a preview image, using the per-file example image the
|
||||
site publishes when the README has none
|
||||
6. Updates the scanner cache
|
||||
7. Broadcasts WebSocket progress events
|
||||
|
||||
#### Site-provided card extras (`fetch_model_card_context`)
|
||||
|
||||
A model card is not always just `README.md`. ModelScope keeps the author's
|
||||
summary (`Description`), the site-curated tags (`OfficialTags`), and — per
|
||||
published version — the model filenames together with that file's example
|
||||
images (`MuseInfo.versions[].coverImages`) and trigger words in its
|
||||
model-detail API. AIGC repositories there often ship an auto-generated
|
||||
boilerplate README and put everything useful in `Description`, so reading only
|
||||
the README yields almost nothing.
|
||||
|
||||
Providers opt in by overriding `ModelSource.fetch_model_card_context()`, which
|
||||
returns a `ModelCardContext`. The wanted file is identified by its sha256 when
|
||||
the caller knows it (the scanner already records one) and by **basename**
|
||||
otherwise, so each checkpoint in a collection repo gets its own images — and
|
||||
keeps getting them after the user renames the weights, which is the only
|
||||
identifier a rename cannot invalidate. Sites with no such extras inherit an
|
||||
empty context, and the pipeline behaves exactly as before.
|
||||
|
||||
The README and the repository metadata describe the whole repository, not one
|
||||
file, so `execute_skill()` creates a `ModelSourceCache` for the duration of a
|
||||
run and passes it down. Enriching the eight checkpoints of one ModelScope
|
||||
repository costs two HTTP requests instead of sixteen; only the per-file
|
||||
selection is redone for each file. Nothing is cached across runs, and download
|
||||
URLs never go through it.
|
||||
|
||||
#### Deterministic data is applied whether or not an LLM is configured
|
||||
|
||||
`AgentService._load_source_card()` runs for every source-backed enrichment, and
|
||||
the post-processor applies what it returns before the LLM output is merged. A
|
||||
user with **no** provider configured therefore still gets the author summary,
|
||||
the example images, the preview, the site-curated tags, the trigger words and
|
||||
the README rendered as the model description.
|
||||
|
||||
The LLM is always consulted when one is configured — invoking **Enrich Metadata
|
||||
with AI** must call the provider every time, and the site data is never treated
|
||||
as a reason to skip it. The deterministic values act as fallbacks that fill
|
||||
gaps the LLM leaves behind:
|
||||
|
||||
| Field | Deterministic source | LLM role |
|
||||
| --- | --- | --- |
|
||||
| `modelDescription` | author summary + README as HTML | — |
|
||||
| `civitai.images` | site example images, then README images | — |
|
||||
| `preview_url` | first available example image | may propose one from the README |
|
||||
| `tags` | site-curated tags, always merged in | proposes additional content tags |
|
||||
| `civitai.description` | author summary | richer 1-2 sentence summary wins |
|
||||
| `base_model` | site hints resolved against the canonical vocabulary (`py/services/agent/base_model_resolver.py`) | mapping it is the LLM's job; the resolver only fills in when the LLM returns nothing |
|
||||
| `trainedWords` | per-file site trigger words, then YAML `instance_prompt` | primary extraction |
|
||||
| `usage_tips` | regex over an explicitly stated strength range | primary extraction |
|
||||
| `notes` | — | LLM-only |
|
||||
|
||||
Models with no source, an unknown source, or a source without model-card access (TensorArt) are skipped with an explicit reason and counted in the run summary.
|
||||
|
||||
**Model types**: LoRA, Checkpoint, Embedding
|
||||
|
||||
@@ -27,11 +27,14 @@ from typing import Any, Dict, List, Optional
|
||||
from ...config import config
|
||||
from ..llm_service import LLMService
|
||||
from ..model_sources import (
|
||||
ModelCardContext,
|
||||
ModelSourceCache,
|
||||
get_source,
|
||||
resolve_source_ref,
|
||||
source_label,
|
||||
)
|
||||
from ..websocket_manager import ws_manager
|
||||
from .base_model_resolver import resolve_base_model
|
||||
from .post_processor import PostProcessor
|
||||
from .skill_registry import SkillRegistry
|
||||
from .skills.enrich_hf_metadata.readme_processor import (
|
||||
@@ -257,6 +260,11 @@ class AgentService:
|
||||
llm = await self._ensure_llm()
|
||||
llm_configured = llm.is_configured() if skill.llm_required else True
|
||||
|
||||
# A collection repository holds many model files under one source id;
|
||||
# this memo keeps the README and the repository metadata from being
|
||||
# re-fetched once per file. It lives for this run only.
|
||||
source_cache = ModelSourceCache()
|
||||
|
||||
for model_path in model_paths:
|
||||
model_filename = os.path.basename(model_path)
|
||||
logger.info(
|
||||
@@ -282,14 +290,37 @@ class AgentService:
|
||||
skip_model = True
|
||||
|
||||
if not skip_model:
|
||||
prompt_vars: Dict[str, Any] = {"model_path": model_path}
|
||||
if skill.llm_required and llm_configured:
|
||||
prompt_vars = await self._build_prompt_context(
|
||||
skill_name, model_path, metadata, registry, llm,
|
||||
# The site's own data is deterministic and must land whether
|
||||
# or not an LLM is available: a user without a key still gets
|
||||
# the author summary, the example images and the tags.
|
||||
source_vars, source_context = await self._load_source_card(
|
||||
model_path, metadata, cache=source_cache,
|
||||
)
|
||||
resolved_base_model = ""
|
||||
if skill_name == "enrich_hf_metadata" and not (
|
||||
metadata.get("base_model") or ""
|
||||
).strip():
|
||||
resolved_base_model = await self._resolve_site_base_model(
|
||||
source_context,
|
||||
)
|
||||
|
||||
llm_response: Optional[Dict[str, Any]] = None
|
||||
if skill.llm_required and llm_configured:
|
||||
if skill.llm_required and not llm_configured:
|
||||
# Without a provider the deterministic model-source data
|
||||
# still lands; the LLM-only fields simply stay untouched.
|
||||
logger.info(
|
||||
"[%s] No LLM configured for %s — applying %s data only",
|
||||
skill_name, model_filename,
|
||||
"model-source"
|
||||
if not source_context.is_empty()
|
||||
else "README",
|
||||
)
|
||||
elif skill.llm_required:
|
||||
prompt_vars = await self._build_prompt_context(
|
||||
skill_name, model_path, metadata, registry, llm,
|
||||
source_vars=source_vars,
|
||||
source_context=source_context,
|
||||
)
|
||||
prompt_template = registry.load_prompt(skill_name)
|
||||
rendered = _render_prompt(prompt_template, prompt_vars)
|
||||
llm_response = await llm.chat_completion_json(
|
||||
@@ -312,7 +343,9 @@ class AgentService:
|
||||
model_path=model_path,
|
||||
llm_output=llm_response or {},
|
||||
metadata=metadata,
|
||||
readme_content=prompt_vars.get("readme_content_full", ""),
|
||||
readme_content=source_vars.get("readme_content_full", ""),
|
||||
source_context=source_context,
|
||||
resolved_base_model=resolved_base_model,
|
||||
)
|
||||
|
||||
if model_result.get("success", True):
|
||||
@@ -395,6 +428,97 @@ class AgentService:
|
||||
"""
|
||||
return "\n".join(f"- {m}" for m in models)
|
||||
|
||||
async def _load_source_card(
|
||||
self,
|
||||
model_path: str,
|
||||
metadata: Dict[str, Any],
|
||||
*,
|
||||
cache: Optional[ModelSourceCache] = None,
|
||||
) -> tuple[Dict[str, Any], ModelCardContext]:
|
||||
"""Fetch the model card and site-published extras for one model.
|
||||
|
||||
Runs for every source-backed enrichment regardless of LLM
|
||||
availability, because everything it returns is deterministic data that
|
||||
should be applied even without a configured provider.
|
||||
|
||||
*cache* is the per-run memo created by :meth:`execute_skill`. The
|
||||
README is repository-wide, so it is fetched once per source id; only
|
||||
successful reads are memoised, leaving a transient failure to be
|
||||
retried for the next file.
|
||||
"""
|
||||
|
||||
variables: Dict[str, Any] = {
|
||||
"asset_base_url": "",
|
||||
"source_description": "",
|
||||
"source_base_model": "",
|
||||
"source_official_tags": "",
|
||||
"source_example_images": "",
|
||||
"source_trigger_words": "",
|
||||
"readme_content": "(README not available)",
|
||||
"readme_content_full": "",
|
||||
}
|
||||
|
||||
ref = resolve_source_ref(metadata)
|
||||
source = get_source(ref.platform) if ref is not None else None
|
||||
if ref is None or source is None or not source.supports_enrichment:
|
||||
return variables, ModelCardContext()
|
||||
|
||||
raw_basename = os.path.splitext(os.path.basename(model_path))[0]
|
||||
variables["asset_base_url"] = source.asset_base_url(ref.source_id)
|
||||
|
||||
cache_key = f"{ref.platform}:{ref.source_id}"
|
||||
readme = cache.readmes.get(cache_key) if cache is not None else None
|
||||
if readme is None:
|
||||
readme = await source.fetch_model_card(ref.source_id)
|
||||
if cache is not None and readme:
|
||||
cache.readmes[cache_key] = readme
|
||||
|
||||
# Sites such as ModelScope keep part of the model card outside the
|
||||
# README (author summary, curated tags, per-file example images). The
|
||||
# recorded hash identifies the file even after the user renames it.
|
||||
card_context = await source.fetch_model_card_context(
|
||||
ref.source_id,
|
||||
os.path.basename(model_path),
|
||||
sha256=(metadata.get("sha256") or "").strip(),
|
||||
cache=cache,
|
||||
)
|
||||
variables["source_description"] = card_context.description
|
||||
variables["source_base_model"] = card_context.base_model
|
||||
variables["source_official_tags"] = "\n".join(
|
||||
f"- {tag}" for tag in card_context.official_tags
|
||||
)
|
||||
variables["source_example_images"] = "\n".join(
|
||||
f"- {url}" for url in card_context.example_images
|
||||
)
|
||||
variables["source_trigger_words"] = ", ".join(card_context.trigger_words)
|
||||
|
||||
# Trim README to the section relevant to this model file
|
||||
# (collection repos often have multiple models in one README).
|
||||
if readme and raw_basename:
|
||||
trimmed = extract_relevant_section(readme, raw_basename)
|
||||
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
|
||||
else:
|
||||
cleaned = clean_readme_for_llm(readme) if readme else ""
|
||||
variables["readme_content"] = cleaned if cleaned else "(README not available)"
|
||||
variables["readme_content_full"] = readme or ""
|
||||
|
||||
return variables, card_context
|
||||
|
||||
async def _resolve_site_base_model(self, source_context: ModelCardContext) -> str:
|
||||
"""Resolve the site's base-model hints to a canonical name, or ``""``."""
|
||||
|
||||
from ...metadata_ops import list_base_models
|
||||
|
||||
hints = [*source_context.base_model_aliases, source_context.base_model]
|
||||
if not any(hints):
|
||||
return ""
|
||||
try:
|
||||
known_names = await list_base_models()
|
||||
except Exception as exc:
|
||||
logger.debug("Failed to list base models for site resolution: %s", exc)
|
||||
return ""
|
||||
return resolve_base_model(hints, known_names)
|
||||
|
||||
async def _build_prompt_context(
|
||||
self,
|
||||
skill_name: str,
|
||||
@@ -402,16 +526,25 @@ class AgentService:
|
||||
metadata: Dict[str, Any],
|
||||
registry: SkillRegistry,
|
||||
llm: Any,
|
||||
*,
|
||||
source_vars: Optional[Dict[str, Any]] = None,
|
||||
source_context: Optional[ModelCardContext] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Gather variables for the skill's prompt template.
|
||||
|
||||
Reads metadata, fetches the HF README (if applicable), lists available
|
||||
Reads metadata, fetches the model card (unless a pre-fetched
|
||||
*source_vars* / *source_context* pair is supplied), lists available
|
||||
base models, loads user priority tags, and returns a dict that maps to
|
||||
``{{variable}}`` placeholders in ``prompt.md``.
|
||||
"""
|
||||
from ...metadata_ops import identify_model_type, list_base_models
|
||||
from ..settings_manager import SettingsManager
|
||||
|
||||
if source_vars is None or source_context is None:
|
||||
source_vars, source_context = await self._load_source_card(
|
||||
model_path, metadata,
|
||||
)
|
||||
|
||||
context: Dict[str, Any] = {
|
||||
"model_path": model_path,
|
||||
"model_basename": "",
|
||||
@@ -421,6 +554,15 @@ class AgentService:
|
||||
"source_platform": "",
|
||||
"source_label": "",
|
||||
"asset_base_url": "",
|
||||
# Site-provided card extras (see ModelSource.fetch_model_card_context)
|
||||
"source_description": "",
|
||||
"source_base_model": "",
|
||||
"source_official_tags": "",
|
||||
"source_example_images": "",
|
||||
"source_trigger_words": "",
|
||||
# Carrier for the structured context handed to the post-processor;
|
||||
# never rendered into the prompt.
|
||||
"source_context": ModelCardContext(),
|
||||
# Legacy Hugging Face aliases (kept so older prompt templates and
|
||||
# third-party skills keep rendering)
|
||||
"hf_url": "",
|
||||
@@ -458,17 +600,17 @@ class AgentService:
|
||||
|
||||
source = get_source(ref.platform) if ref is not None else None
|
||||
if ref is not None and source is not None and source.supports_enrichment:
|
||||
context["asset_base_url"] = source.asset_base_url(ref.source_id)
|
||||
readme = await source.fetch_model_card(ref.source_id)
|
||||
# Trim README to the section relevant to this model file
|
||||
# (collection repos often have multiple models in one README).
|
||||
if readme and raw_basename:
|
||||
trimmed = extract_relevant_section(readme, raw_basename)
|
||||
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
|
||||
else:
|
||||
cleaned = clean_readme_for_llm(readme) if readme else ""
|
||||
context["readme_content"] = cleaned if cleaned else "(README not available)"
|
||||
context["readme_content_full"] = readme or ""
|
||||
# Values fetched once by _load_source_card and shared with the
|
||||
# post-processor, so the network is not hit twice per model.
|
||||
context["asset_base_url"] = source_vars["asset_base_url"]
|
||||
context["source_context"] = source_context
|
||||
context["source_description"] = source_vars["source_description"]
|
||||
context["source_base_model"] = source_vars["source_base_model"]
|
||||
context["source_official_tags"] = source_vars["source_official_tags"]
|
||||
context["source_example_images"] = source_vars["source_example_images"]
|
||||
context["source_trigger_words"] = source_vars["source_trigger_words"]
|
||||
context["readme_content"] = source_vars["readme_content"]
|
||||
context["readme_content_full"] = source_vars["readme_content_full"]
|
||||
|
||||
try:
|
||||
raw_models = await list_base_models()
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Map a site-reported base model onto this system's canonical vocabulary.
|
||||
|
||||
Model sites name base models in their own terms: ModelScope publishes
|
||||
``krea/Krea-2-Turbo`` and ``KREA_2_TURBO`` where this system expects the
|
||||
canonical ``Krea 2``. Turning one into the other is normally the LLM's job;
|
||||
this module resolves the cases that can be decided safely so the canonical
|
||||
field is still populated when the LLM returns nothing usable for it.
|
||||
|
||||
The resolver is deliberately strict, because a wrong base model written with
|
||||
apparent authority is worse than no value at all:
|
||||
|
||||
* it only ever returns a name that is already present in *known_names*;
|
||||
* matching is on the normalised form (lowercased, non-alphanumerics removed),
|
||||
so separators and casing are ignored but nothing is inferred;
|
||||
* a bounded set of published variant suffixes may be stripped, and only when
|
||||
the remainder still matches a known name exactly.
|
||||
|
||||
Anything it cannot decide returns ``""``, and the caller falls back to the LLM.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Iterable, Sequence
|
||||
|
||||
#: Variant suffixes sites append to a base-model *family* name. Stripping one
|
||||
#: is only attempted when the remainder matches a known name exactly, so an
|
||||
#: unrecognised suffix can never produce a bogus match.
|
||||
_VARIANT_SUFFIXES: tuple[str, ...] = (
|
||||
"turbo",
|
||||
"schnell",
|
||||
"lightning",
|
||||
"dev",
|
||||
"beta",
|
||||
"alpha",
|
||||
)
|
||||
|
||||
_NON_ALNUM = re.compile(r"[^a-z0-9]+")
|
||||
|
||||
|
||||
def _normalize(value: str) -> str:
|
||||
"""Return the comparison form of *value*.
|
||||
|
||||
Lowercases and drops every non-alphanumeric character, so ``KREA_2``,
|
||||
``Krea 2``, ``krea-2`` and ``krea.2`` all collapse to ``krea2``.
|
||||
"""
|
||||
|
||||
return _NON_ALNUM.sub("", (value or "").lower())
|
||||
|
||||
|
||||
def resolve_base_model(
|
||||
hints: Iterable[str], known_names: Sequence[str]
|
||||
) -> str:
|
||||
"""Return the canonical base model that *hints* refers to, or ``""``.
|
||||
|
||||
Args:
|
||||
hints: Site-reported names, best first (e.g. an architecture enum
|
||||
before a link-style repository id).
|
||||
known_names: The canonical vocabulary; only these are ever returned.
|
||||
|
||||
Returns:
|
||||
One of *known_names*, or ``""`` when nothing matches exactly.
|
||||
"""
|
||||
|
||||
normalized: dict[str, str] = {}
|
||||
for name in known_names:
|
||||
key = _normalize(name)
|
||||
if key and key not in normalized:
|
||||
normalized[key] = name
|
||||
if not normalized:
|
||||
return ""
|
||||
|
||||
ordered = [hint for hint in hints if hint]
|
||||
|
||||
# 1. Exact normalised match — the unambiguous case.
|
||||
for hint in ordered:
|
||||
candidate = _normalize(hint)
|
||||
if candidate in normalized:
|
||||
return normalized[candidate]
|
||||
|
||||
# 2. Drop one published variant suffix and retry exactly.
|
||||
for hint in ordered:
|
||||
candidate = _normalize(hint)
|
||||
for suffix in _VARIANT_SUFFIXES:
|
||||
if not candidate.endswith(suffix) or candidate == suffix:
|
||||
continue
|
||||
stem = candidate[: -len(suffix)]
|
||||
if stem in normalized:
|
||||
return normalized[stem]
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
__all__ = ["resolve_base_model"]
|
||||
@@ -10,12 +10,16 @@ refresh cache). All actual I/O is delegated to :mod:`~py.metadata_ops`.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from ..model_sources import ModelCardContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,6 +46,8 @@ class PostProcessor:
|
||||
llm_output: Dict[str, Any],
|
||||
metadata: Dict[str, Any],
|
||||
readme_content: str = "",
|
||||
source_context: Optional["ModelCardContext"] = None,
|
||||
resolved_base_model: str = "",
|
||||
) -> Dict[str, Any]:
|
||||
"""Route *llm_output* to the correct skill post-processor.
|
||||
|
||||
@@ -49,12 +55,21 @@ class PostProcessor:
|
||||
that is converted to HTML and stored as ``modelDescription`` for
|
||||
the description tab.
|
||||
|
||||
*source_context* carries the extras the model site publishes outside
|
||||
the README (author description, per-file example images, trigger
|
||||
words). It is ``None`` for callers that have none.
|
||||
|
||||
*resolved_base_model* is the canonical base-model name the site's own
|
||||
hints resolve to, used when the LLM did not supply one (which is the
|
||||
normal case when the LLM was skipped).
|
||||
|
||||
Returns a dict with keys ``success`` (bool), ``updated_fields`` (list),
|
||||
``preview_downloaded`` (bool), and ``errors`` (list).
|
||||
"""
|
||||
if skill_name == "enrich_hf_metadata":
|
||||
return await self._process_enrich_hf_metadata(
|
||||
model_path, llm_output, metadata, readme_content,
|
||||
model_path, llm_output, metadata, readme_content, source_context,
|
||||
resolved_base_model,
|
||||
)
|
||||
return {
|
||||
"success": False,
|
||||
@@ -72,6 +87,8 @@ class PostProcessor:
|
||||
llm_output: Dict[str, Any],
|
||||
metadata: Dict[str, Any],
|
||||
readme_content: str = "",
|
||||
source_context: Optional["ModelCardContext"] = None,
|
||||
resolved_base_model: str = "",
|
||||
) -> Dict[str, Any]:
|
||||
from ...metadata_ops import (
|
||||
apply_metadata_updates,
|
||||
@@ -109,8 +126,11 @@ class PostProcessor:
|
||||
# -- Collect updates -----------------------------------------------
|
||||
updates: Dict[str, Any] = {}
|
||||
|
||||
# base_model
|
||||
# base_model — the LLM's mapping wins; when it returned nothing usable,
|
||||
# fall back to the canonical name the site's own hints resolve to.
|
||||
new_base = (llm_output.get("base_model") or "").strip()
|
||||
if not new_base:
|
||||
new_base = (resolved_base_model or "").strip()
|
||||
current_base = metadata.get("base_model", "") or ""
|
||||
if new_base and self._should_overwrite(current_base, is_source_model):
|
||||
updates["base_model"] = new_base
|
||||
@@ -131,14 +151,29 @@ class PostProcessor:
|
||||
trig_civitai["trainedWords"] = cleaned
|
||||
updates["civitai"] = trig_civitai
|
||||
|
||||
# modelDescription — from raw README content (converted to HTML)
|
||||
if readme_content and is_source_model:
|
||||
converted = convert_readme_to_html(readme_content)
|
||||
if converted:
|
||||
updates["modelDescription"] = converted
|
||||
# modelDescription — the author's own summary (when the site keeps one
|
||||
# outside the README, e.g. ModelScope's ``Description``) followed by the
|
||||
# README converted to HTML.
|
||||
site_description = (
|
||||
(source_context.description if source_context else "") or ""
|
||||
).strip()
|
||||
if is_source_model and (site_description or readme_content):
|
||||
parts: List[str] = []
|
||||
if site_description:
|
||||
parts.append(f"<p>{html.escape(site_description)}</p>")
|
||||
if readme_content:
|
||||
converted = convert_readme_to_html(readme_content)
|
||||
if converted:
|
||||
parts.append(converted)
|
||||
if parts:
|
||||
updates["modelDescription"] = "\n".join(parts)
|
||||
|
||||
# short_description → civitai.description (for "About this version")
|
||||
# short_description → civitai.description (for "About this version").
|
||||
# Falls back to the site's author summary, which for ModelScope AIGC
|
||||
# models is frequently the only human-written text available.
|
||||
short_desc = (llm_output.get("short_description") or "").strip()
|
||||
if not short_desc:
|
||||
short_desc = site_description
|
||||
if short_desc and is_source_model:
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
desc_civitai = dict(current_civitai)
|
||||
@@ -147,19 +182,31 @@ class PostProcessor:
|
||||
desc_civitai["description"] = short_desc
|
||||
updates["civitai"] = desc_civitai
|
||||
|
||||
# gallery images → civitai.images (from YAML frontmatter widget entries
|
||||
# and Sample Gallery markdown tables in the README body)
|
||||
gallery_images: List[Dict[str, Any]] = []
|
||||
if readme_content and is_source_model:
|
||||
repo = source_id
|
||||
if repo:
|
||||
rec_w = llm_output.get("recommended_width") or 0
|
||||
rec_h = llm_output.get("recommended_height") or 0
|
||||
# gallery images → civitai.images (site example images, YAML frontmatter
|
||||
# widget entries, and Sample Gallery markdown tables in the README body)
|
||||
rec_width = llm_output.get("recommended_width") or 0
|
||||
rec_height = llm_output.get("recommended_height") or 0
|
||||
|
||||
# Example images the site publishes for *this* file. They are matched
|
||||
# by filename, so they are the most precise preview source available
|
||||
# and the only one for repositories whose README carries no images.
|
||||
site_images: List[Dict[str, Any]] = []
|
||||
if is_source_model and source_context is not None:
|
||||
site_images = [
|
||||
_example_image(url, rec_width, rec_height)
|
||||
for url in source_context.example_images
|
||||
if url
|
||||
]
|
||||
|
||||
gallery_images: List[Dict[str, Any]] = []
|
||||
if (readme_content or site_images) and is_source_model:
|
||||
repo = source_id
|
||||
readme_images: List[Dict[str, Any]] = []
|
||||
if readme_content and repo:
|
||||
# 1. Widget images (YAML frontmatter)
|
||||
gallery = extract_gallery_images(
|
||||
readme_content, repo,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
default_width=rec_width, default_height=rec_height,
|
||||
base_url=asset_base_url,
|
||||
)
|
||||
|
||||
@@ -168,7 +215,7 @@ class PostProcessor:
|
||||
table_images = extract_gallery_table_images(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
default_width=rec_width, default_height=rec_height,
|
||||
base_url=asset_base_url,
|
||||
)
|
||||
existing_urls.update(img["url"] for img in table_images if img.get("url"))
|
||||
@@ -177,7 +224,7 @@ class PostProcessor:
|
||||
simple_images = extract_simple_markdown_images(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
default_width=rec_width, default_height=rec_height,
|
||||
base_url=asset_base_url,
|
||||
)
|
||||
existing_urls.update(img["url"] for img in simple_images if img.get("url"))
|
||||
@@ -186,25 +233,39 @@ class PostProcessor:
|
||||
html_images = extract_html_img_tags(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
default_width=rec_width, default_height=rec_height,
|
||||
base_url=asset_base_url,
|
||||
)
|
||||
|
||||
all_images = gallery + table_images + simple_images + html_images
|
||||
if all_images:
|
||||
gallery_images = all_images
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
gallery_civitai = dict(current_civitai)
|
||||
if "civitai" in updates and isinstance(updates["civitai"], dict):
|
||||
gallery_civitai.update(updates["civitai"])
|
||||
gallery_civitai["images"] = all_images
|
||||
updates["civitai"] = gallery_civitai
|
||||
readme_images = gallery + table_images + simple_images + html_images
|
||||
|
||||
# tags
|
||||
# Site images come first so the preview fallback below prefers an
|
||||
# image that is known to belong to this exact file.
|
||||
all_images = _dedupe_images(site_images + readme_images)
|
||||
if all_images:
|
||||
gallery_images = all_images
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
gallery_civitai = dict(current_civitai)
|
||||
if "civitai" in updates and isinstance(updates["civitai"], dict):
|
||||
gallery_civitai.update(updates["civitai"])
|
||||
gallery_civitai["images"] = all_images
|
||||
updates["civitai"] = gallery_civitai
|
||||
|
||||
# tags — the site's curated tags are authoritative content vocabulary, so
|
||||
# they are kept alongside whatever the LLM proposed (the LLM is skipped
|
||||
# entirely when the site data is complete, which is why this cannot rely
|
||||
# on ``llm_output`` alone).
|
||||
new_tags = llm_output.get("tags", [])
|
||||
if isinstance(new_tags, list) and new_tags:
|
||||
candidate_tags: List[str] = []
|
||||
if is_source_model and source_context is not None:
|
||||
candidate_tags.extend(source_context.official_tags)
|
||||
if isinstance(new_tags, list):
|
||||
candidate_tags.extend(
|
||||
tag for tag in new_tags if tag not in candidate_tags
|
||||
)
|
||||
if candidate_tags:
|
||||
existing_tags = metadata.get("tags") or []
|
||||
merged = self._merge_tags(existing_tags, new_tags)
|
||||
merged = self._merge_tags(existing_tags, candidate_tags)
|
||||
if len(merged) > len(existing_tags) or is_source_model:
|
||||
updates["tags"] = merged
|
||||
|
||||
@@ -212,21 +273,30 @@ class PostProcessor:
|
||||
updates["metadata_source"] = "agent:enrich_hf_metadata"
|
||||
updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Store LLM confidence in metadata so it's accessible for evaluation
|
||||
# LLM confidence, stored for the enrichment evaluation harness. The key
|
||||
# must NOT start with an underscore: `BaseModelMetadata.from_dict()`
|
||||
# deliberately drops underscore-prefixed keys so they never round-trip,
|
||||
# which silently erased this field on the next metadata write.
|
||||
raw_confidence = (llm_output.get("confidence") or "").strip()
|
||||
if raw_confidence:
|
||||
updates["_llm_confidence"] = raw_confidence
|
||||
updates["llm_confidence"] = raw_confidence
|
||||
|
||||
# Fallback: extract instance_prompt from YAML frontmatter when the LLM
|
||||
# returned empty trigger words but the README has instance_prompt.
|
||||
# Fallback: use the trigger words the site records for this exact file,
|
||||
# then the README's YAML `instance_prompt`, when the LLM returned none.
|
||||
if trigger_words_empty:
|
||||
instance_prompt = _extract_yaml_instance_prompt(readme_content)
|
||||
if instance_prompt:
|
||||
site_triggers = (
|
||||
list(source_context.trigger_words) if source_context else []
|
||||
)
|
||||
if not site_triggers:
|
||||
instance_prompt = _extract_yaml_instance_prompt(readme_content)
|
||||
if instance_prompt:
|
||||
site_triggers = [instance_prompt]
|
||||
if site_triggers:
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
trig_civitai = dict(current_civitai)
|
||||
if "civitai" in updates and isinstance(updates["civitai"], dict):
|
||||
trig_civitai.update(updates["civitai"])
|
||||
trig_civitai["trainedWords"] = [instance_prompt]
|
||||
trig_civitai["trainedWords"] = site_triggers
|
||||
updates["civitai"] = trig_civitai
|
||||
|
||||
preview_remote_url = (llm_output.get("preview_url") or "").strip()
|
||||
@@ -260,8 +330,12 @@ class PostProcessor:
|
||||
if new_notes:
|
||||
updates["notes"] = new_notes
|
||||
|
||||
# usage_tips — JSON string (e.g. {"strength_min":0.85,"strength_max":1.4})
|
||||
# usage_tips — JSON string (e.g. {"strength_min":0.85,"strength_max":1.4}).
|
||||
# When the LLM returned nothing, recover an explicitly stated strength
|
||||
# range from the author summary so the value is not lost.
|
||||
raw_tips = (llm_output.get("usage_tips") or "").strip()
|
||||
if not raw_tips or raw_tips == "{}":
|
||||
raw_tips = _extract_usage_tips(site_description)
|
||||
if raw_tips and raw_tips != "{}":
|
||||
try:
|
||||
json.loads(raw_tips)
|
||||
@@ -324,6 +398,129 @@ class PostProcessor:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
#: Separator between a label and its value. Published model cards routinely
|
||||
#: wrap the numbers in markdown emphasis or quotes (``strength: **0.85 - 1.4**``,
|
||||
#: ``CLIP 强度「0.5」``), so those are absorbed rather than treated as a break.
|
||||
_EMPHASIS = "[\"'\u201c\u201d\u300c\u300d*_`\\s]*"
|
||||
|
||||
#: An explicitly stated strength/weight range, e.g. ``权重0.5-1.2``,
|
||||
#: ``强度 0.8 ~ 1.2``, ``strength: **0.85 - 1.4**``.
|
||||
_RANGE_DASH = "(?:-|\u2010|\u2011|\u2012|\u2013|\u2014|\uff0d|~|\uff5e|\u81f3|\u5230|to)"
|
||||
|
||||
_STRENGTH_RANGE_RE = re.compile(
|
||||
"(?:\u6743\u91cd|\u5f3a\u5ea6|strength|weight)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
|
||||
+ r"(\d+(?:\.\d+)?)" + _EMPHASIS + _RANGE_DASH + _EMPHASIS
|
||||
+ r"(\d+(?:\.\d+)?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
#: A single strength/weight value, e.g. ``strength: 0.6``, ``权重 0.8``.
|
||||
_STRENGTH_VALUE_RE = re.compile(
|
||||
"(?:\u6743\u91cd|\u5f3a\u5ea6|strength|weight)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
|
||||
+ r"(\d+(?:\.\d+)?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
#: ``clip strength: 0.5`` / ``CLIP 强度 0.5``.
|
||||
_CLIP_STRENGTH_RE = re.compile(
|
||||
"clip" + _EMPHASIS + "(?:\u5f3a\u5ea6|strength)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
|
||||
+ r"(\d+(?:\.\d+)?)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
#: ``clip skip: 2`` / ``CLIP 跳过 2``.
|
||||
_CLIP_SKIP_RE = re.compile(
|
||||
"clip" + _EMPHASIS + "(?:skip|\u8df3\u8fc7)" + _EMPHASIS + "[:\uff1a]?" + _EMPHASIS
|
||||
+ r"(\d+)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _extract_usage_tips(text: str) -> str:
|
||||
"""Extract stated strength/CLIP recommendations from prose.
|
||||
|
||||
This is the deterministic counterpart to the LLM's ``usage_tips`` output,
|
||||
used when the LLM was skipped. It only recognises explicitly written
|
||||
values — it never infers a range — and returns ``""`` when it finds none.
|
||||
|
||||
Returns:
|
||||
A JSON string matching the skill's ``usage_tips`` schema, or ``""``.
|
||||
"""
|
||||
|
||||
if not text:
|
||||
return ""
|
||||
|
||||
tips: Dict[str, Any] = {}
|
||||
|
||||
# CLIP strength is resolved first and then blanked out, so the generic
|
||||
# strength patterns cannot mistake `CLIP 强度 0.5` for the LoRA strength.
|
||||
text_for_strength = text
|
||||
clip_strength = _CLIP_STRENGTH_RE.search(text_for_strength)
|
||||
if clip_strength:
|
||||
tips["clip_strength"] = float(clip_strength.group(1))
|
||||
text_for_strength = (
|
||||
text_for_strength[: clip_strength.start()]
|
||||
+ " "
|
||||
+ text_for_strength[clip_strength.end() :]
|
||||
)
|
||||
|
||||
range_match = _STRENGTH_RANGE_RE.search(text_for_strength)
|
||||
if range_match:
|
||||
low = float(range_match.group(1))
|
||||
high = float(range_match.group(2))
|
||||
if low > high:
|
||||
low, high = high, low
|
||||
tips["strength_min"] = low
|
||||
tips["strength_max"] = high
|
||||
tips["strength_range"] = f"{low:g}-{high:g}"
|
||||
else:
|
||||
value_match = _STRENGTH_VALUE_RE.search(text_for_strength)
|
||||
if value_match:
|
||||
tips["strength"] = float(value_match.group(1))
|
||||
|
||||
clip_skip = _CLIP_SKIP_RE.search(text)
|
||||
if clip_skip:
|
||||
tips["clip_skip"] = int(clip_skip.group(1))
|
||||
|
||||
if not tips:
|
||||
return ""
|
||||
return json.dumps(tips, ensure_ascii=False)
|
||||
|
||||
|
||||
def _example_image(url: str, width: int, height: int) -> Dict[str, Any]:
|
||||
"""Build a ``civitai.images`` entry for a site-provided example image.
|
||||
|
||||
The site publishes no prompt alongside these images, so the entry carries
|
||||
empty prompt metadata and the LLM's recommended dimensions when it found
|
||||
any (falling back to the same 512px placeholder the README extractors use).
|
||||
"""
|
||||
|
||||
return {
|
||||
"url": url,
|
||||
"type": "image",
|
||||
"nsfwLevel": 0,
|
||||
"width": width or 512,
|
||||
"height": height or 512,
|
||||
"meta": {"prompt": "", "negativePrompt": ""},
|
||||
"hasMeta": False,
|
||||
"hasPositivePrompt": False,
|
||||
}
|
||||
|
||||
|
||||
def _dedupe_images(images: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Drop later entries that repeat an earlier image URL, keeping order."""
|
||||
|
||||
seen: set[str] = set()
|
||||
unique: List[Dict[str, Any]] = []
|
||||
for image in images:
|
||||
url = image.get("url") or ""
|
||||
if not url or url in seen:
|
||||
continue
|
||||
seen.add(url)
|
||||
unique.append(image)
|
||||
return unique
|
||||
|
||||
|
||||
def _extract_yaml_instance_prompt(readme_content: str) -> str:
|
||||
"""Extract ``instance_prompt`` from the YAML frontmatter of a HF README.
|
||||
|
||||
|
||||
@@ -25,6 +25,34 @@ You are an expert assistant for AI image generation models. Your task is to extr
|
||||
{{current_metadata}}
|
||||
```
|
||||
|
||||
## Site-Provided Metadata (any field may be empty)
|
||||
|
||||
The model site publishes the following **alongside** the README. It is
|
||||
first-hand information recorded by the site itself, so it outranks anything
|
||||
you would otherwise guess:
|
||||
|
||||
- **Author description**: {{source_description}}
|
||||
- **Base model reported by the site**: {{source_base_model}}
|
||||
- **Trigger words recorded for this file**: {{source_trigger_words}}
|
||||
- **Site-curated tags**:
|
||||
{{source_official_tags}}
|
||||
- **Example image URLs for this file**:
|
||||
{{source_example_images}}
|
||||
|
||||
Use it as follows:
|
||||
|
||||
- A weight or strength range stated in the **author description** belongs in
|
||||
``usage_tips`` (and in ``notes``); do not leave ``usage_tips`` empty when the
|
||||
description states one.
|
||||
- When the author description exists, base ``short_description`` on it rather
|
||||
than on the README, which on some sites is auto-generated boilerplate.
|
||||
- Treat the **site-curated tags** as strong signals for ``tags``: they are
|
||||
already a curated content vocabulary, so prefer them over invented words.
|
||||
- Treat the **base model reported by the site** as a strong hint for
|
||||
``base_model``, but still map it to the EXACT canonical name from the
|
||||
available base-model list.
|
||||
- Use the **example image URLs** when the README contains no usable image.
|
||||
|
||||
## User Priority Tags Reference
|
||||
|
||||
The user has configured the following list of **meaningful tag categories** for this model type (`{{model_type}}`):
|
||||
@@ -55,10 +83,11 @@ Extract the following information from the README content above:
|
||||
### base_model
|
||||
The base model this model was trained on. Use EXACTLY one of the names from the **Available Base Models** list above. Do not invent new names or use aliases.
|
||||
|
||||
Check the YAML frontmatter for ``base_model:`` first. If the frontmatter has no ``base_model:``, look at the **model filename** (``{{model_basename}}``), YAML ``tags:``, README title and first paragraph for clues — the base model family is often embedded in the name
|
||||
Check the **base model reported by the site** (above) and the YAML frontmatter ``base_model:`` first. If neither yields a match, look at the **model filename** (``{{model_basename}}``), YAML ``tags:``, README title and first paragraph for clues — the base model family is often embedded in the name
|
||||
|
||||
### trigger_words
|
||||
The trigger words or activation prompts needed to use this LoRA. Look for:
|
||||
- The **trigger words recorded for this file** in the site-provided metadata (most authoritative)
|
||||
- `instance_prompt:` in the YAML frontmatter
|
||||
- Phrases like "trigger word:", "trigger:", "use this prompt:", "activation prompt:"
|
||||
- In collection repos: the trigger section **specific to this model file** (look near matching download links or anchor IDs)
|
||||
@@ -66,12 +95,13 @@ The trigger words or activation prompts needed to use this LoRA. Look for:
|
||||
Return as an array of strings. If none found, return an empty array `[]`. **Never** return `["None"]` or any placeholder value — a truly empty list means no trigger words exist.
|
||||
|
||||
### short_description
|
||||
A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`, not the repo as a whole. Return empty string if the README is too minimal.
|
||||
A concise 1-2 sentence summary of what this model does. For collection repos, focus on the **specific model version** matching `{{model_basename}}`, not the repo as a whole. Prefer the **author description** from the site-provided metadata when it is present; otherwise extract from the "Model description" section or the first paragraph. Return empty string if the available content is too minimal.
|
||||
|
||||
### tags
|
||||
3-8 relevant tags for categorizing this model. **Quality over quantity.**
|
||||
|
||||
Sources to consider:
|
||||
- The **site-curated tags** from the site-provided metadata (these are already filtered content tags — prefer them)
|
||||
- The YAML frontmatter `tags:` list (filter out technical ones — see below)
|
||||
- The subject, style, character, or concept the model represents
|
||||
- The model filename itself may give clues (e.g. "pokemon", "anime", "pixelart")
|
||||
@@ -82,7 +112,9 @@ Sources to consider:
|
||||
|
||||
2. **Cross-reference against the priority_tags reference.** Only include a tag if it meaningfully describes what the model actually creates (subject, style, character type) and is semantically close to one of the priority_tags. If none of the README's tags match meaningful categories, prefer returning a smaller set or an empty array over including low-value tags.
|
||||
|
||||
3. **All lowercase, no spaces, no hyphens** (use single words like `"photorealistic"`, `"anime"`, `"character"`).
|
||||
3. **All lowercase, and keep each tag's own wording.** Prefer the spelling already used by the site, the frontmatter, or the author — including hyphenated and multi-word tags such as `"sci-fi"`, `"semi-realistic"`, `"character-enhancement"` or `"art style"`. Do **not** strip separators or invent a single-word variant of a tag you are already including (e.g. do not emit both `"character-enhancement"` and `"character"`). When a tag is written in another script (e.g. Chinese), likewise keep it verbatim instead of translating it.
|
||||
|
||||
4. **Never invent a tag** that neither the site-provided metadata, the YAML frontmatter, nor the README text supports.
|
||||
|
||||
Return empty array if no meaningful content tags remain after filtering.
|
||||
|
||||
@@ -95,13 +127,13 @@ The URL of the most suitable preview image from the README. Look for:
|
||||
- The YAML frontmatter `widget:` section (which often has `output.url` fields)
|
||||
- In collection repos: the sample images listed **under the section** for this specific model version
|
||||
- Generic `` in the body
|
||||
Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL from the repository raw-file base URL (`{{asset_base_url}}`) plus the relative path. If no suitable image is found, return an empty string.
|
||||
Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL from the repository raw-file base URL (`{{asset_base_url}}`) plus the relative path. If the README has no suitable image, fall back to the site-provided **example image URLs** for this file. If nothing is available, return an empty string.
|
||||
|
||||
### notes
|
||||
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`. Return empty string if the README has no useful usage info.
|
||||
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`. Include the **author description** from the site-provided metadata when it is present. Return empty string if there is no useful usage info.
|
||||
|
||||
### usage_tips
|
||||
A JSON string with structured usage recommendations. Extract from the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5"). Possible fields (include only those you can determine):
|
||||
A JSON string with structured usage recommendations. Extract from the **author description** (site-provided metadata) and the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5", "权重0.5-1.2"). Possible fields (include only those you can determine):
|
||||
|
||||
```json
|
||||
{
|
||||
|
||||
@@ -392,12 +392,18 @@ def _extract_frontmatter(text: str) -> str:
|
||||
|
||||
|
||||
def convert_readme_to_html(markdown_text: str | None) -> str:
|
||||
"""Convert HF README markdown to sanitised HTML."""
|
||||
"""Convert HF README markdown to sanitised HTML.
|
||||
|
||||
Site-generated placeholder notices are dropped here too, so a repository
|
||||
whose author wrote nothing does not store the download instructions as its
|
||||
model description; the result is an empty string in that case.
|
||||
"""
|
||||
if not markdown_text:
|
||||
return ""
|
||||
|
||||
text = markdown_text
|
||||
text = _strip_frontmatter(text)
|
||||
text = _strip_generated_card_boilerplate(text)
|
||||
text = _strip_gallery(text)
|
||||
text = _strip_badge_images(text)
|
||||
text = _strip_html_comments(text)
|
||||
@@ -444,6 +450,59 @@ _MASSIVE_LIST_LINE_MIN_LEN = 150
|
||||
#: Minimum consecutive enumeration lines to trigger massive-list stripping.
|
||||
_MASSIVE_LIST_THRESHOLD = 8
|
||||
|
||||
#: Substrings identifying text a *site* generated to fill a model card whose
|
||||
#: author wrote nothing, as opposed to the author's own content. ModelScope
|
||||
#: renders such a card as a placeholder notice, a block of SDK/git download
|
||||
#: instructions, and a closing invitation to improve the card.
|
||||
#:
|
||||
#: Matched as substrings rather than whole headings because the notices are
|
||||
#: prose, and because non-Latin scripts are not space-delimited — the notice
|
||||
#: continues with a full-width period, so the ``title == kw`` style matching
|
||||
#: used for :data:`_BOILERPLATE_HEADERS` would never fire.
|
||||
_GENERATED_CARD_MARKERS: tuple[str, ...] = (
|
||||
"当前模型的贡献者未提供更加详细的模型介绍",
|
||||
"您可以通过如下",
|
||||
"如果您是本模型的贡献者",
|
||||
)
|
||||
|
||||
|
||||
def _strip_generated_card_boilerplate(text: str) -> str:
|
||||
"""Remove the notices a site generates to fill an empty model card.
|
||||
|
||||
A repository whose uploader wrote no README still gets a card: ModelScope
|
||||
answers with "the contributor provided no further description", the SDK
|
||||
and git download commands, and an invitation to complete the card. None
|
||||
of it describes the model, yet it was landing in both the LLM prompt and
|
||||
the stored description.
|
||||
|
||||
A notice that is a heading takes its whole section with it, so the
|
||||
download block goes too; a stand-alone notice line is dropped on its own.
|
||||
Content the author added later — under a heading of equal or higher
|
||||
level — is kept, so an improved card is not thrown away.
|
||||
"""
|
||||
|
||||
lines = text.split("\n")
|
||||
out: list[str] = []
|
||||
skip_until_level: int | None = None
|
||||
|
||||
for line in lines:
|
||||
level = _heading_level(line)
|
||||
|
||||
if any(marker in line for marker in _GENERATED_CARD_MARKERS):
|
||||
if level > 0:
|
||||
skip_until_level = level
|
||||
continue
|
||||
|
||||
if skip_until_level is not None:
|
||||
if level > 0 and level <= skip_until_level:
|
||||
skip_until_level = None
|
||||
else:
|
||||
continue
|
||||
|
||||
out.append(line)
|
||||
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def clean_readme_for_llm(markdown_text: str | None, max_length: int = 6000) -> str:
|
||||
"""Clean a HF README for injection into an LLM metadata-extraction prompt.
|
||||
@@ -453,6 +512,8 @@ def clean_readme_for_llm(markdown_text: str | None, max_length: int = 6000) -> s
|
||||
|
||||
* ``widget:`` YAML block (example prompts + output URLs)
|
||||
* ``<Gallery />`` tags and wrappers
|
||||
* Site-generated placeholder notices for a card the author never wrote
|
||||
(see :func:`_strip_generated_card_boilerplate`)
|
||||
* Fenced code blocks (Python / bash / bibtex / yaml)
|
||||
* Standalone ```` image lines and ``<img>`` tags
|
||||
* Training-parameter tables
|
||||
@@ -478,6 +539,7 @@ def clean_readme_for_llm(markdown_text: str | None, max_length: int = 6000) -> s
|
||||
# Order matters — broader strips first, then finer ones.
|
||||
text = _strip_gallery(text)
|
||||
text = _strip_widget_section(text)
|
||||
text = _strip_generated_card_boilerplate(text)
|
||||
text = _strip_fenced_code_blocks(text)
|
||||
text = _strip_standalone_images(text)
|
||||
text = _strip_training_tables(text)
|
||||
|
||||
@@ -11,7 +11,9 @@ from __future__ import annotations
|
||||
from .base import (
|
||||
GROUP_PREFIXES,
|
||||
HTTP_TIMEOUT,
|
||||
ModelCardContext,
|
||||
ModelSource,
|
||||
ModelSourceCache,
|
||||
ModelSourceError,
|
||||
SourceRef,
|
||||
USER_AGENT,
|
||||
@@ -45,7 +47,9 @@ __all__ = [
|
||||
"GROUP_PREFIXES",
|
||||
"HTTP_TIMEOUT",
|
||||
"LEGACY_HF_URL_FIELD",
|
||||
"ModelCardContext",
|
||||
"ModelSource",
|
||||
"ModelSourceCache",
|
||||
"ModelSourceError",
|
||||
"HuggingFaceSource",
|
||||
"ModelScopeSource",
|
||||
|
||||
@@ -8,6 +8,8 @@ know about such a site is expressed by :class:`ModelSource`:
|
||||
* how to recognise one of its URLs (:meth:`ModelSource.parse`)
|
||||
* the canonical page URL for a source id (:meth:`ModelSource.canonical_url`)
|
||||
* how to fetch the model card (:meth:`ModelSource.fetch_model_card`)
|
||||
* how to fetch the extras that live *outside* the README
|
||||
(:meth:`ModelSource.fetch_model_card_context`)
|
||||
* how to turn repository-relative asset paths into absolute URLs
|
||||
(:meth:`ModelSource.asset_base_url`)
|
||||
* which capabilities the site actually supports
|
||||
@@ -22,8 +24,8 @@ from __future__ import annotations
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, Optional
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Iterable, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
@@ -61,6 +63,56 @@ class SourceRef:
|
||||
"""Canonical URL of the model page."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelCardContext:
|
||||
"""Site-specific extras that accompany a model's README model card.
|
||||
|
||||
A model card is not always just ``README.md``. ModelScope, for example,
|
||||
keeps the author's summary, the site-curated tags, and the per-file
|
||||
example images in its model-detail API rather than in the repository.
|
||||
Sources with no such extras return an empty context (the default), so
|
||||
every field here must be treated as optional by callers.
|
||||
"""
|
||||
|
||||
description: str = ""
|
||||
"""Author-written summary shown on the model page, outside the README."""
|
||||
|
||||
base_model: str = ""
|
||||
"""Base model as reported by the site (possibly a site-local id)."""
|
||||
|
||||
base_model_aliases: list[str] = field(default_factory=list)
|
||||
"""Other names the site uses for the same base model.
|
||||
|
||||
Sites often publish both a link-style id (``krea/Krea-2-Turbo``) and an
|
||||
internal architecture enum (``KREA_2``). The enum usually normalises
|
||||
cleanly onto this system's canonical vocabulary, so it is the better
|
||||
resolution hint for :mod:`py.services.agent.base_model_resolver`.
|
||||
"""
|
||||
|
||||
official_tags: list[str] = field(default_factory=list)
|
||||
"""Content tags curated by the site itself."""
|
||||
|
||||
example_images: list[str] = field(default_factory=list)
|
||||
"""Absolute URLs of example images for the requested model file."""
|
||||
|
||||
trigger_words: list[str] = field(default_factory=list)
|
||||
"""Trigger words the site records for the requested model file."""
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""Return ``True`` when the site contributed nothing extra."""
|
||||
|
||||
return not any(
|
||||
(
|
||||
self.description,
|
||||
self.base_model,
|
||||
self.base_model_aliases,
|
||||
self.official_tags,
|
||||
self.example_images,
|
||||
self.trigger_words,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class ModelSourceError(Exception):
|
||||
"""Raised when a model source cannot satisfy a request.
|
||||
|
||||
@@ -73,6 +125,26 @@ class ModelSourceError(Exception):
|
||||
self.status = status
|
||||
|
||||
|
||||
class ModelSourceCache:
|
||||
"""Per-run memo shared between the agent pipeline and a model source.
|
||||
|
||||
A collection repository publishes many model files under a single source
|
||||
id, so enriching each file re-fetches the same README and the same
|
||||
repository metadata. One cache is created per enrichment run and thrown
|
||||
away afterwards: nothing is retained across runs (a model card can change
|
||||
at any time), and download URLs are never routed through it.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
#: Provider-agnostic: ``"<platform>:<source_id>"`` → raw README text.
|
||||
self.readmes: Dict[str, str] = {}
|
||||
#: Provider-owned scratch space. Keys must be namespaced by the
|
||||
#: provider (``(platform, kind, source_id)``) so two providers can
|
||||
#: never collide. Only successful results should be stored, so a
|
||||
#: transient failure is still retried for the next file.
|
||||
self.provider: Dict[Any, Any] = {}
|
||||
|
||||
|
||||
#: Repository ids are always exactly ``owner/name``. Components may contain
|
||||
#: dots (``black-forest-labs/FLUX.1-dev``) but must not be empty, ``.`` / ``..``,
|
||||
#: or start with a dot - the id is used as a path segment on disk.
|
||||
@@ -230,6 +302,35 @@ class ModelSource:
|
||||
|
||||
return ""
|
||||
|
||||
async def fetch_model_card_context(
|
||||
self,
|
||||
source_id: str,
|
||||
filename: str = "",
|
||||
*,
|
||||
sha256: str = "",
|
||||
cache: Optional["ModelSourceCache"] = None,
|
||||
) -> ModelCardContext:
|
||||
"""Return the card extras the site keeps outside the README.
|
||||
|
||||
*filename* is the model file's basename (no directory) and *sha256*
|
||||
its content hash; between them they select the right entry when a
|
||||
repository holds several models. A site that records per-file hashes
|
||||
should prefer *sha256*, because it is the only identifier that
|
||||
survives the user renaming the weights.
|
||||
|
||||
*cache* is an optional per-run memo (see :class:`ModelSourceCache`)
|
||||
that lets a provider avoid re-fetching repository-wide data for every
|
||||
file in a collection repository.
|
||||
|
||||
Sites whose model card is fully described by :meth:`fetch_model_card`
|
||||
need no override and inherit this empty context.
|
||||
|
||||
Implementations must never raise: enrichment treats a missing
|
||||
context as "the site had nothing extra to say".
|
||||
"""
|
||||
|
||||
return ModelCardContext()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Download support
|
||||
# ------------------------------------------------------------------
|
||||
@@ -301,7 +402,9 @@ def filter_weight_files(entries: Iterable[tuple[str, int]]) -> list[dict[str, An
|
||||
__all__ = [
|
||||
"GROUP_PREFIXES",
|
||||
"HTTP_TIMEOUT",
|
||||
"ModelCardContext",
|
||||
"ModelSource",
|
||||
"ModelSourceCache",
|
||||
"ModelSourceError",
|
||||
"SourceRef",
|
||||
"USER_AGENT",
|
||||
|
||||
@@ -2,13 +2,19 @@
|
||||
|
||||
ModelScope exposes the same "model card as README.md" convention as
|
||||
Hugging Face, including a YAML frontmatter block that often carries
|
||||
``base_model:`` and ``trigger_words:``. Three public endpoints are used,
|
||||
``base_model:`` and ``trigger_words:``. Four public endpoints are used,
|
||||
none of which requires an API key for public models:
|
||||
|
||||
* ``/models/{owner}/{name}/resolve/{revision}/README.md`` — raw model card
|
||||
* ``/api/v1/models/{owner}/{name}/repo?Revision=..&FilePath=README.md`` —
|
||||
the same content through the API, used as a fallback when the resolve
|
||||
URL is unavailable.
|
||||
* ``/api/v1/models/{owner}/{name}`` — the model-detail payload behind the
|
||||
model page. It carries the author's summary (``Description``), the
|
||||
site-curated tags (``OfficialTags``), and, per published version, the
|
||||
model filenames (``MuseInfo.versions[].stats.fileList``) together with
|
||||
that file's example images (``coverImages``) and trigger words. See
|
||||
:meth:`ModelScopeSource.fetch_model_card_context`.
|
||||
* ``/api/v1/models/{owner}/{name}/repo/files?Revision=..`` — the file
|
||||
listing backing the download picker. It reports real sizes for LFS
|
||||
files (not the pointer size), so no extra HEAD request is needed.
|
||||
@@ -18,14 +24,22 @@ which redirects to a CDN URL carrying a time-limited ``auth_key``.
|
||||
Requesting the resolve URL fresh on every attempt (which the shared
|
||||
downloader does, including for resumable Range requests) keeps that key
|
||||
valid; the CDN URL must never be cached.
|
||||
|
||||
The README and the detail payload both describe the whole repository rather
|
||||
than one file, so a per-run ``ModelSourceCache`` keeps them from being read
|
||||
again for every checkpoint of a collection repository.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from .base import (
|
||||
ModelCardContext,
|
||||
ModelSource,
|
||||
ModelSourceError,
|
||||
fetch_json,
|
||||
@@ -33,6 +47,9 @@ from .base import (
|
||||
filter_weight_files,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from .base import ModelSourceCache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_URL_PATTERN = re.compile(
|
||||
@@ -95,6 +112,67 @@ class ModelScopeSource(ModelSource):
|
||||
return text
|
||||
return ""
|
||||
|
||||
async def fetch_model_card_context(
|
||||
self,
|
||||
source_id: str,
|
||||
filename: str = "",
|
||||
*,
|
||||
sha256: str = "",
|
||||
cache: Optional["ModelSourceCache"] = None,
|
||||
) -> ModelCardContext:
|
||||
"""Read the model-detail API that backs the ModelScope model page.
|
||||
|
||||
ModelScope splits a model card in two: ``README.md`` holds the
|
||||
long-form content, while the author's summary, the site-curated tags,
|
||||
and the per-file example images live only here. AIGC repositories
|
||||
frequently ship an auto-generated README ("the contributor provided
|
||||
no further description") and put everything useful in ``Description``,
|
||||
so enrichment that reads only the README comes back nearly empty.
|
||||
|
||||
The wanted file is identified by its sha256 when the caller knows it
|
||||
and by *filename* otherwise; see :func:`_matching_versions`. The
|
||||
images and trigger words returned belong to that exact
|
||||
``.safetensors`` — essential for collection repositories, where every
|
||||
checkpoint has its own sample image.
|
||||
|
||||
The detail payload describes the whole repository and is therefore
|
||||
shared across every file in it, so it is read through *cache* when the
|
||||
caller supplies one; only the per-file selection is redone.
|
||||
"""
|
||||
|
||||
data = await self._fetch_detail(source_id, cache=cache)
|
||||
if data is None:
|
||||
return ModelCardContext()
|
||||
return _build_card_context(data, filename, sha256)
|
||||
|
||||
async def _fetch_detail(
|
||||
self,
|
||||
source_id: str,
|
||||
*,
|
||||
cache: Optional["ModelSourceCache"] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Fetch (or reuse) the model-detail payload for *source_id*."""
|
||||
|
||||
cache_key = (self.platform, "detail", source_id)
|
||||
if cache is not None and cache_key in cache.provider:
|
||||
return cache.provider[cache_key]
|
||||
|
||||
status, payload = await fetch_json(
|
||||
f"https://modelscope.cn/api/v1/models/{source_id}"
|
||||
)
|
||||
if status != 200 or not isinstance(payload, dict):
|
||||
logger.debug(
|
||||
"ModelScope detail API returned HTTP %s for %s", status, source_id
|
||||
)
|
||||
return None
|
||||
data = payload.get("Data")
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
|
||||
if cache is not None:
|
||||
cache.provider[cache_key] = data
|
||||
return data
|
||||
|
||||
async def list_files(
|
||||
self, source_id: str, revision: str = ""
|
||||
) -> list[dict]:
|
||||
@@ -142,3 +220,294 @@ class ModelScopeSource(ModelSource):
|
||||
|
||||
|
||||
__all__ = ["ModelScopeSource"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model-detail API parsing helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Trigger-word values that mean "the author left this blank".
|
||||
_EMPTY_TRIGGER_VALUES = frozenset({"none", "null", "n/a"})
|
||||
|
||||
|
||||
def _clean_text(value: Any) -> str:
|
||||
"""Return a stripped string for *value*, or ``""`` for anything else."""
|
||||
|
||||
return value.strip() if isinstance(value, str) else ""
|
||||
|
||||
|
||||
def _first_string(value: Any) -> str:
|
||||
"""Return the first non-empty string in a list, or ``""``."""
|
||||
|
||||
if isinstance(value, list):
|
||||
for item in value:
|
||||
text = _clean_text(item)
|
||||
if text:
|
||||
return text
|
||||
return ""
|
||||
|
||||
|
||||
def _build_card_context(
|
||||
data: dict[str, Any], filename: str, sha256: str = ""
|
||||
) -> ModelCardContext:
|
||||
"""Turn a model-detail payload into a :class:`ModelCardContext`.
|
||||
|
||||
Separated from the HTTP fetch so the repository-wide payload can be cached
|
||||
across the files of a collection repository while the per-file selection
|
||||
is still redone for each one.
|
||||
"""
|
||||
|
||||
context = ModelCardContext(
|
||||
description=_clean_text(data.get("Description")),
|
||||
base_model=_first_string(data.get("BaseModel")),
|
||||
base_model_aliases=_base_model_aliases(data),
|
||||
official_tags=_official_tags(data.get("OfficialTags")),
|
||||
)
|
||||
|
||||
versions = _matching_versions(
|
||||
data.get("MuseInfo"),
|
||||
filename,
|
||||
digests=_file_digests(data),
|
||||
sha256=sha256,
|
||||
)
|
||||
if versions:
|
||||
context.example_images = _cover_image_urls(versions)
|
||||
context.trigger_words = _version_trigger_words(versions)
|
||||
return context
|
||||
|
||||
|
||||
def _base_model_aliases(data: dict[str, Any]) -> list[str]:
|
||||
"""Return the site's own names for the base model.
|
||||
|
||||
ModelScope publishes a link-style id (``krea/Krea-2-Turbo``) plus its
|
||||
internal architecture enums (``VisionFoundation: KREA_2``,
|
||||
``SubVisionFoundation: KREA_2_TURBO``). The enums are the better
|
||||
resolution hint because they normalise onto this system's canonical
|
||||
vocabulary, so they come first; the owner prefix is also stripped from
|
||||
the link-style ids.
|
||||
"""
|
||||
|
||||
aliases: list[str] = []
|
||||
for key in ("VisionFoundation", "SubVisionFoundation"):
|
||||
value = _clean_text(data.get(key))
|
||||
if value and value not in aliases:
|
||||
aliases.append(value)
|
||||
|
||||
base_models = data.get("BaseModel")
|
||||
if isinstance(base_models, list):
|
||||
for item in base_models:
|
||||
text = _clean_text(item)
|
||||
leaf = text.rsplit("/", 1)[-1] if text else ""
|
||||
if leaf and leaf not in aliases:
|
||||
aliases.append(leaf)
|
||||
return aliases
|
||||
|
||||
|
||||
def _official_tags(value: Any) -> list[str]:
|
||||
"""Extract the site-curated tag values from ``OfficialTags``.
|
||||
|
||||
ModelScope's entries are dicts carrying an English ``Tag`` plus a
|
||||
``ChineseName``; the English value is the curated content vocabulary, so
|
||||
that is the one surfaced here.
|
||||
"""
|
||||
|
||||
tags: list[str] = []
|
||||
if not isinstance(value, list):
|
||||
return tags
|
||||
for entry in value:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
tag = _clean_text(entry.get("Tag"))
|
||||
if tag and tag not in tags:
|
||||
tags.append(tag)
|
||||
return tags
|
||||
|
||||
|
||||
def _version_files(version: dict[str, Any]) -> list[str]:
|
||||
"""Return the model filenames covered by one ``MuseInfo.versions`` entry.
|
||||
|
||||
The listing normally sits in ``stats.fileList``; some payloads only
|
||||
carry the same field as a JSON-encoded string under
|
||||
``modelVersion.stats``, so both shapes are accepted.
|
||||
"""
|
||||
|
||||
stats = version.get("stats")
|
||||
files = stats.get("fileList") if isinstance(stats, dict) else None
|
||||
|
||||
if not isinstance(files, list):
|
||||
model_version = version.get("modelVersion")
|
||||
raw = model_version.get("stats") if isinstance(model_version, dict) else None
|
||||
if isinstance(raw, str) and raw.strip():
|
||||
try:
|
||||
decoded = json.loads(raw)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
decoded = None
|
||||
if isinstance(decoded, dict):
|
||||
files = decoded.get("fileList")
|
||||
|
||||
if not isinstance(files, list):
|
||||
return []
|
||||
return [item for item in files if isinstance(item, str) and item]
|
||||
|
||||
|
||||
def _version_show_name(version: dict[str, Any]) -> str:
|
||||
"""Return the human-facing version label (e.g. ``c1-st1000``)."""
|
||||
|
||||
model_version = version.get("modelVersion")
|
||||
if not isinstance(model_version, dict):
|
||||
return ""
|
||||
return _clean_text(model_version.get("showName")).lower()
|
||||
|
||||
|
||||
def _file_digests(data: dict[str, Any]) -> dict[str, str]:
|
||||
"""Return ``basename -> sha256`` for every published weight file.
|
||||
|
||||
``ModelInfos`` groups the repository's files by kind (``safetensor``,
|
||||
…) and records a real sha256 for each, which is what makes it possible to
|
||||
recognise a file the user has renamed.
|
||||
"""
|
||||
|
||||
digests: dict[str, str] = {}
|
||||
model_infos = data.get("ModelInfos")
|
||||
if not isinstance(model_infos, dict):
|
||||
return digests
|
||||
for info in model_infos.values():
|
||||
files = info.get("files") if isinstance(info, dict) else None
|
||||
if not isinstance(files, list):
|
||||
continue
|
||||
for entry in files:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
name = _clean_text(entry.get("name"))
|
||||
digest = _clean_text(entry.get("sha256"))
|
||||
if name and digest:
|
||||
digests.setdefault(os.path.basename(name).lower(), digest.lower())
|
||||
return digests
|
||||
|
||||
|
||||
def _matching_versions(
|
||||
muse_info: Any,
|
||||
filename: str,
|
||||
*,
|
||||
digests: dict[str, str] | None = None,
|
||||
sha256: str = "",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return the ``versions`` entries that publish the wanted model file.
|
||||
|
||||
Strategies, in order:
|
||||
|
||||
1. **sha256** — the file's content hash, looked up through
|
||||
:func:`_file_digests`. This is the only strategy that survives the
|
||||
user renaming the weights, which is common once a model is filed away.
|
||||
2. **Exact basename** against each version's ``stats.fileList``.
|
||||
3. **``showName`` inside the file stem**, which absorbs the naming drift
|
||||
ModelScope sometimes applies to uploaded weights.
|
||||
|
||||
A known-but-unmatched hash falls through to the filename strategies
|
||||
rather than giving up, in case the local file was re-encoded. All matches
|
||||
are returned so a file re-published across several versions contributes
|
||||
all of its example images. With no *filename* and no *sha256*, only an
|
||||
unambiguous single-version repository is used, because a per-file image
|
||||
must never be attributed to the wrong file.
|
||||
"""
|
||||
|
||||
if not isinstance(muse_info, dict):
|
||||
return []
|
||||
versions = muse_info.get("versions")
|
||||
if not isinstance(versions, list):
|
||||
return []
|
||||
entries = [entry for entry in versions if isinstance(entry, dict)]
|
||||
if not entries:
|
||||
return []
|
||||
|
||||
target_hash = (sha256 or "").strip().lower()
|
||||
if target_hash:
|
||||
known = digests or {}
|
||||
by_hash: list[dict[str, Any]] = []
|
||||
for version in entries:
|
||||
for path in _version_files(version):
|
||||
if known.get(os.path.basename(path).lower()) == target_hash:
|
||||
by_hash.append(version)
|
||||
break
|
||||
if by_hash:
|
||||
return by_hash
|
||||
|
||||
if not filename:
|
||||
return entries if len(entries) == 1 else []
|
||||
|
||||
target = os.path.basename(filename).strip().lower()
|
||||
if not target:
|
||||
return []
|
||||
stem = os.path.splitext(target)[0]
|
||||
|
||||
exact: list[dict[str, Any]] = []
|
||||
fuzzy: list[dict[str, Any]] = []
|
||||
for version in entries:
|
||||
files = {os.path.basename(path).lower() for path in _version_files(version)}
|
||||
if target in files:
|
||||
exact.append(version)
|
||||
continue
|
||||
show_name = _version_show_name(version)
|
||||
if show_name and show_name in stem:
|
||||
fuzzy.append(version)
|
||||
|
||||
return exact or fuzzy
|
||||
|
||||
|
||||
def _cover_image_urls(versions: list[dict[str, Any]]) -> list[str]:
|
||||
"""Collect the example-image URLs published by the given versions."""
|
||||
|
||||
urls: list[str] = []
|
||||
for version in versions:
|
||||
covers = version.get("coverImages")
|
||||
if not isinstance(covers, list):
|
||||
continue
|
||||
for cover in covers:
|
||||
if not isinstance(cover, dict):
|
||||
continue
|
||||
url = _clean_text(cover.get("url"))
|
||||
if url and url not in urls:
|
||||
urls.append(url)
|
||||
return urls
|
||||
|
||||
|
||||
def _version_trigger_words(versions: list[dict[str, Any]]) -> list[str]:
|
||||
"""Return the first non-empty trigger-word list across *versions*."""
|
||||
|
||||
for version in versions:
|
||||
model_version = version.get("modelVersion")
|
||||
raw = (
|
||||
model_version.get("triggerWords")
|
||||
if isinstance(model_version, dict)
|
||||
else None
|
||||
)
|
||||
words = _parse_trigger_words(raw)
|
||||
if words:
|
||||
return words
|
||||
return []
|
||||
|
||||
|
||||
def _parse_trigger_words(raw: Any) -> list[str]:
|
||||
"""Decode ModelScope's JSON-encoded trigger-word string list."""
|
||||
|
||||
if isinstance(raw, list):
|
||||
candidates = raw
|
||||
elif isinstance(raw, str) and raw.strip():
|
||||
try:
|
||||
decoded = json.loads(raw)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return []
|
||||
if not isinstance(decoded, list):
|
||||
return []
|
||||
candidates = decoded
|
||||
else:
|
||||
return []
|
||||
|
||||
words: list[str] = []
|
||||
for item in candidates:
|
||||
word = _clean_text(item)
|
||||
if not word or word.lower() in _EMPTY_TRIGGER_VALUES:
|
||||
continue
|
||||
if word not in words:
|
||||
words.append(word)
|
||||
return words
|
||||
|
||||
@@ -750,6 +750,32 @@
|
||||
#downloadModal .modal-content {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
overflow: hidden; /* The active step scrolls instead of the whole modal */
|
||||
}
|
||||
|
||||
/* Sticky footer layout (mirrors the import modal fix): fixed header,
|
||||
scrollable step content, pinned action buttons. Ensures Back/Download
|
||||
buttons stay visible on short viewports (e.g. 1080p or 150% zoom). */
|
||||
#downloadModal .modal-header {
|
||||
flex-shrink: 0;
|
||||
}
|
||||
|
||||
#downloadModal .download-step {
|
||||
flex: 1 1 auto;
|
||||
min-height: 0; /* Allow the step to shrink and scroll within the flex container */
|
||||
overflow-y: auto;
|
||||
overflow-x: hidden;
|
||||
scrollbar-gutter: stable;
|
||||
}
|
||||
|
||||
#downloadModal .download-step .modal-actions {
|
||||
position: sticky;
|
||||
bottom: 0;
|
||||
z-index: 1;
|
||||
background: var(--lora-surface);
|
||||
border-top: 1px solid var(--lora-border);
|
||||
padding-top: var(--space-2);
|
||||
padding-bottom: var(--space-1);
|
||||
}
|
||||
|
||||
#batchPreviewStep {
|
||||
|
||||
@@ -108,7 +108,12 @@ def evaluate_model(
|
||||
model_description: str = metadata.get("modelDescription") or ""
|
||||
base_model: str = metadata.get("base_model") or ""
|
||||
preview_url: str = metadata.get("preview_url") or ""
|
||||
confidence: str = metadata.get("_llm_confidence") or ""
|
||||
# `_llm_confidence` is the legacy key: underscore-prefixed metadata keys are
|
||||
# deliberately not persisted through `BaseModelMetadata`, so older sidecars
|
||||
# may still carry it while current ones use `llm_confidence`.
|
||||
confidence: str = (
|
||||
metadata.get("llm_confidence") or metadata.get("_llm_confidence") or ""
|
||||
)
|
||||
|
||||
# --- base_model ---
|
||||
base_model_valid = base_model in SUPPORTED_BASE_MODELS
|
||||
|
||||
@@ -490,3 +490,118 @@ class TestStripFencedCodeBlocks:
|
||||
def test_pattern(self, R):
|
||||
text = "x\n```yaml\nkey: val\n```\ny"
|
||||
assert "key: val" not in R._strip_fenced_code_blocks(text)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Site-generated placeholder cards
|
||||
# ======================================================================
|
||||
|
||||
#: The card ModelScope renders when the uploader wrote no README. Copied from
|
||||
#: a live repository so the marker strings stay honest.
|
||||
PLACEHOLDER_CARD = """---
|
||||
base_model: krea/Krea-2-Turbo
|
||||
license: Apache License 2.0
|
||||
tags:
|
||||
- LoRA
|
||||
- text-to-image
|
||||
- \u68a6\u5e7b\u5149\u5f71
|
||||
---
|
||||
### \u5f53\u524d\u6a21\u578b\u7684\u8d21\u732e\u8005\u672a\u63d0\u4f9b\u66f4\u52a0\u8be6\u7ec6\u7684\u6a21\u578b\u4ecb\u7ecd\u3002\u6a21\u578b\u6587\u4ef6\u548c\u6743\u91cd\uff0c\u53ef\u6d4f\u89c8\u201c\u6a21\u578b\u6587\u4ef6\u201d\u9875\u9762\u83b7\u53d6\u3002
|
||||
#### \u60a8\u53ef\u4ee5\u901a\u8fc7\u5982\u4e0bgit clone\u547d\u4ee4\uff0c\u6216\u8005ModelScope SDK\u6765\u4e0b\u8f7d\u6a21\u578b
|
||||
|
||||
SDK\u4e0b\u8f7d
|
||||
```bash
|
||||
#\u5b89\u88c5ModelScope
|
||||
pip install modelscope
|
||||
```
|
||||
Git\u4e0b\u8f7d
|
||||
```
|
||||
#Git\u6a21\u578b\u4e0b\u8f7d
|
||||
git clone https://www.modelscope.cn/yan303145427/krea2-CcFQWZ-Portrait.git
|
||||
```
|
||||
|
||||
<p style="color: lightgrey;">\u5982\u679c\u60a8\u662f\u672c\u6a21\u578b\u7684\u8d21\u732e\u8005\uff0c\u6211\u4eec\u9080\u8bf7\u60a8\u6839\u636e<a href="x">\u6a21\u578b\u8d21\u732e\u6587\u6863</a>\uff0c\u53ca\u65f6\u5b8c\u5584\u6a21\u578b\u5361\u7247\u5185\u5bb9\u3002</p>
|
||||
"""
|
||||
|
||||
#: A real, author-written card (ModelScope AIGC training output).
|
||||
REAL_CARD = """---
|
||||
base_model: krea/Krea-2-Turbo
|
||||
---
|
||||
# krea\u8138\u6a21
|
||||
|
||||
## \u6a21\u578b\u4ecb\u7ecd
|
||||
|
||||
\u672c\u6a21\u578b\u4f9d\u6258\u9b54\u642d\u793e\u533a\u5b8c\u6210\u8bad\u7ec3\u3002
|
||||
|
||||
## \u63a8\u7406\u4ee3\u7801
|
||||
|
||||
\u5b89\u88c5 DiffSynth-Studio\uff1a
|
||||
"""
|
||||
|
||||
|
||||
class TestStripGeneratedCardBoilerplate:
|
||||
def test_removes_the_whole_placeholder_body(self, R):
|
||||
stripped = R._strip_generated_card_boilerplate(PLACEHOLDER_CARD)
|
||||
assert "\u5f53\u524d\u6a21\u578b\u7684\u8d21\u732e\u8005" not in stripped
|
||||
assert "SDK\u4e0b\u8f7d" not in stripped
|
||||
assert "git clone" not in stripped
|
||||
assert "\u9080\u8bf7\u60a8" not in stripped
|
||||
# The frontmatter is the only thing that survives.
|
||||
assert "base_model: krea/Krea-2-Turbo" in stripped
|
||||
|
||||
def test_leaves_a_real_card_untouched(self, R):
|
||||
assert R._strip_generated_card_boilerplate(REAL_CARD) == REAL_CARD
|
||||
|
||||
def test_drops_a_standalone_invitation_line(self, R):
|
||||
text = "real body\n<p>\u5982\u679c\u60a8\u662f\u672c\u6a21\u578b\u7684\u8d21\u732e\u8005\uff0c\u8bf7\u5b8c\u5584</p>\nmore body"
|
||||
stripped = R._strip_generated_card_boilerplate(text)
|
||||
assert "real body" in stripped
|
||||
assert "more body" in stripped
|
||||
assert "\u8d21\u732e\u8005" not in stripped
|
||||
|
||||
def test_keeps_content_added_after_the_placeholder(self, R):
|
||||
"""An author who later wrote a real section must not lose it."""
|
||||
text = (
|
||||
"### \u5f53\u524d\u6a21\u578b\u7684\u8d21\u732e\u8005\u672a\u63d0\u4f9b\u66f4\u52a0\u8be6\u7ec6\u7684\u6a21\u578b\u4ecb\u7ecd\u3002\n"
|
||||
"#### \u60a8\u53ef\u4ee5\u901a\u8fc7\u5982\u4e0bgit clone\u547d\u4ee4\u4e0b\u8f7d\u6a21\u578b\n"
|
||||
"```\ngit clone x\n```\n"
|
||||
"## \u6211\u7684\u771f\u5b9e\u4ecb\u7ecd\n"
|
||||
"\u8fd9\u662f\u4f5c\u8005\u540e\u6765\u8865\u5199\u7684\u5185\u5bb9\u3002\n"
|
||||
)
|
||||
stripped = R._strip_generated_card_boilerplate(text)
|
||||
assert "\u6211\u7684\u771f\u5b9e\u4ecb\u7ecd" in stripped
|
||||
assert "\u4f5c\u8005\u540e\u6765\u8865\u5199\u7684\u5185\u5bb9" in stripped
|
||||
assert "git clone" not in stripped
|
||||
|
||||
def test_handles_html_headings(self, R):
|
||||
text = (
|
||||
"<h3>\u5f53\u524d\u6a21\u578b\u7684\u8d21\u732e\u8005\u672a\u63d0\u4f9b\u66f4\u52a0\u8be6\u7ec6\u7684\u6a21\u578b\u4ecb\u7ecd\u3002</h3>\n"
|
||||
"<p>pip install modelscope</p>\n"
|
||||
)
|
||||
assert R._strip_generated_card_boilerplate(text).strip() == ""
|
||||
|
||||
|
||||
class TestCleanReadmeForLlmPlaceholder:
|
||||
def test_boilerplate_is_gone_but_frontmatter_survives(self, R):
|
||||
cleaned = R.clean_readme_for_llm(PLACEHOLDER_CARD)
|
||||
assert "pip install modelscope" not in cleaned
|
||||
assert "git clone" not in cleaned
|
||||
assert "\u5f53\u524d\u6a21\u578b\u7684\u8d21\u732e\u8005" not in cleaned
|
||||
# Metadata the LLM still needs.
|
||||
assert "base_model: krea/Krea-2-Turbo" in cleaned
|
||||
assert "\u68a6\u5e7b\u5149\u5f71" in cleaned
|
||||
|
||||
def test_a_real_card_keeps_its_body(self, R):
|
||||
cleaned = R.clean_readme_for_llm(REAL_CARD)
|
||||
assert "krea\u8138\u6a21" in cleaned
|
||||
assert "\u6a21\u578b\u4ecb\u7ecd" in cleaned
|
||||
|
||||
|
||||
class TestConvertReadmeToHtmlPlaceholder:
|
||||
def test_a_placeholder_card_renders_to_nothing(self, R):
|
||||
assert R.convert_readme_to_html(PLACEHOLDER_CARD) == ""
|
||||
|
||||
def test_a_real_card_still_renders(self, R):
|
||||
html = R.convert_readme_to_html(REAL_CARD)
|
||||
assert "<h1>krea\u8138\u6a21</h1>" in html
|
||||
assert "DiffSynth-Studio" in html
|
||||
|
||||
@@ -11,6 +11,7 @@ from unittest import mock
|
||||
import pytest
|
||||
|
||||
from py.services.agent.agent_service import AgentService
|
||||
from py.services.model_sources import ModelCardContext
|
||||
|
||||
|
||||
class TestEnrichmentSkipReason:
|
||||
@@ -59,12 +60,23 @@ class TestBuildPromptContext:
|
||||
async def test_modelscope_card_populates_source_variables(self):
|
||||
service = AgentService()
|
||||
readme = "---\nbase_model: krea/Krea-2-Turbo\n---\n# krea\n"
|
||||
card_context = ModelCardContext(
|
||||
description="权重0.5-1.2。配合《krea2-Cc-MJ-风格滤镜》lora一起使用。",
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
official_tags=["photography", "woman"],
|
||||
example_images=["https://resources.modelscope.cn/cover-images/a.png"],
|
||||
trigger_words=["kreamodel"],
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value=readme),
|
||||
) as mock_fetch,
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card_context",
|
||||
new=mock.AsyncMock(return_value=card_context),
|
||||
) as mock_context,
|
||||
mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=["Krea 2 Turbo"]),
|
||||
@@ -91,6 +103,12 @@ class TestBuildPromptContext:
|
||||
)
|
||||
|
||||
mock_fetch.assert_awaited_once_with("jj3550945163/Krea-2-LORA")
|
||||
# The per-file lookup must receive the basename, not the full path.
|
||||
mock_context.assert_awaited_once()
|
||||
assert mock_context.call_args.args == (
|
||||
"jj3550945163/Krea-2-LORA",
|
||||
"krea.safetensors",
|
||||
)
|
||||
assert context["source_platform"] == "modelscope"
|
||||
assert context["source_id"] == "jj3550945163/Krea-2-LORA"
|
||||
assert context["source_label"] == "ModelScope"
|
||||
@@ -99,10 +117,57 @@ class TestBuildPromptContext:
|
||||
== "https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master"
|
||||
)
|
||||
assert readme in context["readme_content_full"]
|
||||
# Site-provided extras are rendered into their own prompt variables.
|
||||
assert context["source_description"] == card_context.description
|
||||
assert context["source_base_model"] == "krea/Krea-2-Turbo"
|
||||
assert context["source_official_tags"] == "- photography\n- woman"
|
||||
assert (
|
||||
context["source_example_images"]
|
||||
== "- https://resources.modelscope.cn/cover-images/a.png"
|
||||
)
|
||||
assert context["source_trigger_words"] == "kreamodel"
|
||||
# The structured context is carried through for the post-processor.
|
||||
assert context["source_context"] is card_context
|
||||
# Hugging Face aliases stay empty for a non-HF source.
|
||||
assert context["hf_url"] == ""
|
||||
assert context["repo"] == "jj3550945163/Krea-2-LORA"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_huggingface_card_context_is_empty(self):
|
||||
"""Sources without card extras contribute empty prompt variables."""
|
||||
service = AgentService()
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value="# card\n"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=[]),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.identify_model_type",
|
||||
new=mock.AsyncMock(return_value="lora"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.settings_manager.SettingsManager.get_priority_tag_config",
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
context = await service._build_prompt_context(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/models/loras/thing.safetensors",
|
||||
metadata={"hf_url": "https://huggingface.co/user/repo"},
|
||||
registry=mock.Mock(),
|
||||
llm=mock.Mock(),
|
||||
)
|
||||
|
||||
assert context["source_description"] == ""
|
||||
assert context["source_official_tags"] == ""
|
||||
assert context["source_example_images"] == ""
|
||||
assert context["source_context"].is_empty()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_huggingface_keeps_legacy_aliases(self):
|
||||
service = AgentService()
|
||||
@@ -179,4 +244,325 @@ class TestBuildPromptContext:
|
||||
hf_fetch.assert_not_awaited()
|
||||
ms_fetch.assert_not_awaited()
|
||||
assert context["readme_content"] == ""
|
||||
assert context["source_platform"] == "tensorart"
|
||||
assert context["source_platform"] == "tensorart"
|
||||
|
||||
class TestExecuteSkillThreadsSourceContext:
|
||||
"""The structured site context must reach the post-processor intact."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_source_context_reaches_the_post_processor(self):
|
||||
service = AgentService()
|
||||
card_context = ModelCardContext(
|
||||
example_images=["https://resources.modelscope.cn/cover-images/a.png"]
|
||||
)
|
||||
skill = mock.Mock(llm_required=True, input_schema={})
|
||||
registry = mock.Mock()
|
||||
registry.get_skill.return_value = skill
|
||||
registry.load_prompt.return_value = "{{model_path}}"
|
||||
|
||||
llm = mock.Mock()
|
||||
llm.is_configured.return_value = True
|
||||
llm.chat_completion_json = mock.AsyncMock(return_value={"base_model": "Krea 2"})
|
||||
|
||||
source_vars = {
|
||||
"readme_content_full": "# card",
|
||||
"source_description": "",
|
||||
"source_base_model": "",
|
||||
"source_official_tags": "",
|
||||
"source_example_images": "",
|
||||
"source_trigger_words": "",
|
||||
"asset_base_url": "",
|
||||
"readme_content": "# card",
|
||||
}
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
|
||||
),
|
||||
mock.patch.object(
|
||||
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
|
||||
),
|
||||
mock.patch.object(
|
||||
service,
|
||||
"_load_source_card",
|
||||
new=mock.AsyncMock(return_value=(source_vars, card_context)),
|
||||
),
|
||||
mock.patch.object(
|
||||
service,
|
||||
"_build_prompt_context",
|
||||
new=mock.AsyncMock(
|
||||
return_value={
|
||||
"model_path": "/p.safetensors",
|
||||
"readme_content_full": "# card",
|
||||
"source_context": card_context,
|
||||
}
|
||||
),
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.read_metadata",
|
||||
new=mock.AsyncMock(
|
||||
return_value={
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/u/r",
|
||||
}
|
||||
),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.agent.agent_service.PostProcessor.process",
|
||||
new=mock.AsyncMock(
|
||||
return_value={"success": True, "updated_fields": []}
|
||||
),
|
||||
) as mock_process,
|
||||
):
|
||||
result = await service.execute_skill(
|
||||
skill_name="enrich_hf_metadata",
|
||||
input_data={"model_paths": ["/p.safetensors"]},
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
assert mock_process.call_args.kwargs["source_context"] is card_context
|
||||
assert mock_process.call_args.kwargs["readme_content"] == "# card"
|
||||
|
||||
|
||||
class TestSiteDataAppliedWithoutLlm:
|
||||
"""A: the site's deterministic data must land even with no LLM available."""
|
||||
|
||||
@staticmethod
|
||||
def _run(service, *, llm_configured: bool, card_context: ModelCardContext):
|
||||
skill = mock.Mock(llm_required=True, input_schema={})
|
||||
registry = mock.Mock()
|
||||
registry.get_skill.return_value = skill
|
||||
registry.load_prompt.return_value = "{{model_path}}"
|
||||
|
||||
llm = mock.Mock()
|
||||
llm.is_configured.return_value = llm_configured
|
||||
llm.chat_completion_json = mock.AsyncMock(return_value={"base_model": "Krea 2"})
|
||||
|
||||
source_vars = {
|
||||
"readme_content_full": "# card",
|
||||
"source_description": card_context.description,
|
||||
"source_base_model": card_context.base_model,
|
||||
"source_official_tags": "",
|
||||
"source_example_images": "",
|
||||
"source_trigger_words": "",
|
||||
"asset_base_url": "",
|
||||
"readme_content": "# card",
|
||||
}
|
||||
return skill, registry, llm, source_vars
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unconfigured_llm_still_applies_site_data(self):
|
||||
service = AgentService()
|
||||
card_context = ModelCardContext(
|
||||
description="作者说明",
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
official_tags=["photography"],
|
||||
example_images=["https://cdn.example/a.png"],
|
||||
)
|
||||
skill, registry, llm, source_vars = self._run(
|
||||
service, llm_configured=False, card_context=card_context
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
|
||||
),
|
||||
mock.patch.object(
|
||||
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
|
||||
),
|
||||
mock.patch.object(
|
||||
service,
|
||||
"_load_source_card",
|
||||
new=mock.AsyncMock(return_value=(source_vars, card_context)),
|
||||
),
|
||||
mock.patch.object(
|
||||
service, "_resolve_site_base_model", new=mock.AsyncMock(return_value="Krea 2")
|
||||
),
|
||||
mock.patch(
|
||||
"py.metadata_ops.read_metadata",
|
||||
new=mock.AsyncMock(
|
||||
return_value={
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/u/r",
|
||||
}
|
||||
),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.agent.agent_service.PostProcessor.process",
|
||||
new=mock.AsyncMock(
|
||||
return_value={"success": True, "updated_fields": []}
|
||||
),
|
||||
) as mock_process,
|
||||
):
|
||||
result = await service.execute_skill(
|
||||
skill_name="enrich_hf_metadata",
|
||||
input_data={"model_paths": ["/p.safetensors"]},
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
# No LLM call, but the deterministic payload reached the post-processor.
|
||||
llm.chat_completion_json.assert_not_awaited()
|
||||
kwargs = mock_process.call_args.kwargs
|
||||
assert kwargs["source_context"] is card_context
|
||||
assert kwargs["readme_content"] == "# card"
|
||||
assert kwargs["resolved_base_model"] == "Krea 2"
|
||||
assert kwargs["llm_output"] == {}
|
||||
|
||||
|
||||
class TestLoadSourceCard:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_variables_without_a_source(self):
|
||||
service = AgentService()
|
||||
variables, context = await service._load_source_card("/p.safetensors", {})
|
||||
assert context.is_empty() is True
|
||||
assert variables["readme_content_full"] == ""
|
||||
assert variables["readme_content"] == "(README not available)"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collects_readme_and_site_extras(self):
|
||||
service = AgentService()
|
||||
card = ModelCardContext(
|
||||
description="作者说明",
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
official_tags=["photography", "woman"],
|
||||
example_images=["https://cdn.example/a.png"],
|
||||
trigger_words=["kreaface"],
|
||||
)
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value="# card"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card_context",
|
||||
new=mock.AsyncMock(return_value=card),
|
||||
) as mock_ctx,
|
||||
):
|
||||
variables, context = await service._load_source_card(
|
||||
"/models/loras/krea.safetensors",
|
||||
{
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/u/r",
|
||||
},
|
||||
)
|
||||
|
||||
mock_ctx.assert_awaited_once()
|
||||
assert mock_ctx.call_args.args == ("u/r", "krea.safetensors")
|
||||
assert context is card
|
||||
assert variables["readme_content_full"] == "# card"
|
||||
assert variables["source_description"] == "作者说明"
|
||||
assert variables["source_official_tags"] == "- photography\n- woman"
|
||||
assert variables["source_example_images"] == "- https://cdn.example/a.png"
|
||||
assert variables["source_trigger_words"] == "kreaface"
|
||||
assert variables["asset_base_url"].endswith("/resolve/master")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_site_base_model_uses_the_canonical_vocabulary(self):
|
||||
service = AgentService()
|
||||
with mock.patch(
|
||||
"py.metadata_ops.list_base_models",
|
||||
new=mock.AsyncMock(return_value=["Krea 2", "Flux.1 D"]),
|
||||
):
|
||||
resolved = await service._resolve_site_base_model(
|
||||
ModelCardContext(
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
base_model_aliases=["KREA_2", "KREA_2_TURBO"],
|
||||
)
|
||||
)
|
||||
assert resolved == "Krea 2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_site_base_model_is_empty_without_hints(self):
|
||||
service = AgentService()
|
||||
assert await service._resolve_site_base_model(ModelCardContext()) == ""
|
||||
|
||||
|
||||
class TestLlmAlwaysRunsWhenConfigured:
|
||||
"""Clicking "Enrich Metadata with AI" must always consult the LLM.
|
||||
|
||||
The site-provided data is applied deterministically either way, but it is
|
||||
never treated as a reason to skip the call — the LLM's summary and notes
|
||||
are richer than the raw site fields, and silently not calling out to the
|
||||
provider would make the menu action unpredictable.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_runs_even_when_the_site_supplies_everything(self):
|
||||
service = AgentService()
|
||||
card_context = ModelCardContext(
|
||||
description="作者说明",
|
||||
base_model="krea/Krea-2-Turbo",
|
||||
base_model_aliases=["KREA_2"],
|
||||
official_tags=["photography", "woman"],
|
||||
example_images=["https://cdn.example/a.png"],
|
||||
trigger_words=["kreaface"],
|
||||
)
|
||||
skill = mock.Mock(llm_required=True, input_schema={})
|
||||
registry = mock.Mock()
|
||||
registry.get_skill.return_value = skill
|
||||
registry.load_prompt.return_value = "{{model_path}}"
|
||||
|
||||
llm = mock.Mock()
|
||||
llm.is_configured.return_value = True
|
||||
llm.chat_completion_json = mock.AsyncMock(
|
||||
return_value={"base_model": "Krea 2", "short_description": "llm summary"}
|
||||
)
|
||||
|
||||
source_vars = {
|
||||
"readme_content_full": "# card",
|
||||
"source_description": card_context.description,
|
||||
"source_base_model": card_context.base_model,
|
||||
"source_official_tags": "- photography\n- woman",
|
||||
"source_example_images": "- https://cdn.example/a.png",
|
||||
"source_trigger_words": "kreaface",
|
||||
"asset_base_url": "",
|
||||
"readme_content": "# card",
|
||||
}
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
service, "_ensure_registry", new=mock.AsyncMock(return_value=registry)
|
||||
),
|
||||
mock.patch.object(
|
||||
service, "_ensure_llm", new=mock.AsyncMock(return_value=llm)
|
||||
),
|
||||
mock.patch.object(
|
||||
service,
|
||||
"_load_source_card",
|
||||
new=mock.AsyncMock(return_value=(source_vars, card_context)),
|
||||
),
|
||||
mock.patch.object(
|
||||
service,
|
||||
"_build_prompt_context",
|
||||
new=mock.AsyncMock(
|
||||
return_value={"model_path": "/p.safetensors", "system_prompt": "sys"}
|
||||
),
|
||||
) as mock_prompt,
|
||||
mock.patch(
|
||||
"py.metadata_ops.read_metadata",
|
||||
new=mock.AsyncMock(
|
||||
return_value={
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/u/r",
|
||||
"base_model": "Krea 2",
|
||||
}
|
||||
),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.agent.agent_service.PostProcessor.process",
|
||||
new=mock.AsyncMock(
|
||||
return_value={"success": True, "updated_fields": []}
|
||||
),
|
||||
),
|
||||
):
|
||||
result = await service.execute_skill(
|
||||
skill_name="enrich_hf_metadata",
|
||||
input_data={"model_paths": ["/p.safetensors"]},
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
llm.chat_completion_json.assert_awaited_once()
|
||||
# The prompt is built from the already-fetched card, not re-fetched.
|
||||
assert mock_prompt.call_args.kwargs["source_context"] is card_context
|
||||
assert mock_prompt.call_args.kwargs["source_vars"] is source_vars
|
||||
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Tests for the deterministic site base-model resolver.
|
||||
|
||||
The resolver exists so the enrichment pipeline can skip the LLM when the model
|
||||
site already supplies everything; it must therefore be strictly conservative —
|
||||
returning nothing is always better than returning the wrong canonical name.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.agent.base_model_resolver import resolve_base_model
|
||||
|
||||
KNOWN = [
|
||||
"Krea 2",
|
||||
"Flux.1 Krea",
|
||||
"Flux.1 D",
|
||||
"Flux.1 S",
|
||||
"SDXL 1.0",
|
||||
"Pony",
|
||||
"Illustrious",
|
||||
]
|
||||
|
||||
|
||||
class TestExactNormalisedMatch:
|
||||
@pytest.mark.parametrize(
|
||||
"hint",
|
||||
["Krea 2", "krea 2", "KREA_2", "krea-2", "krea.2", " Krea2 "],
|
||||
)
|
||||
def test_separators_and_casing_are_ignored(self, hint):
|
||||
assert resolve_base_model([hint], KNOWN) == "Krea 2"
|
||||
|
||||
def test_returns_the_canonical_spelling_not_the_hint(self):
|
||||
assert resolve_base_model(["kreA_2"], KNOWN) == "Krea 2"
|
||||
|
||||
def test_does_not_match_a_longer_prefixed_name_by_accident(self):
|
||||
# "Flux.1 Krea" must not be resolved to "Krea 2".
|
||||
assert resolve_base_model(["Flux.1 Krea"], KNOWN) == "Flux.1 Krea"
|
||||
|
||||
def test_first_matching_hint_wins(self):
|
||||
assert (
|
||||
resolve_base_model(["totally-unknown", "KREA_2"], KNOWN) == "Krea 2"
|
||||
)
|
||||
|
||||
|
||||
class TestVariantSuffixStripping:
|
||||
@pytest.mark.parametrize(
|
||||
"hint",
|
||||
[
|
||||
"KREA_2_TURBO",
|
||||
"Krea-2-Turbo",
|
||||
"krea2turbo",
|
||||
"Krea 2 Turbo",
|
||||
"krea-2-dev",
|
||||
"krea2-schnell",
|
||||
"krea2-lightning",
|
||||
],
|
||||
)
|
||||
def test_common_published_suffixes_are_stripped(self, hint):
|
||||
assert resolve_base_model([hint], KNOWN) == "Krea 2"
|
||||
|
||||
def test_suffix_only_hint_never_matches(self):
|
||||
# "turbo" on its own strips to nothing and must not resolve.
|
||||
assert resolve_base_model(["turbo"], KNOWN) == ""
|
||||
|
||||
|
||||
class TestConservativeFailures:
|
||||
@pytest.mark.parametrize(
|
||||
"hints",
|
||||
[
|
||||
[],
|
||||
[""],
|
||||
["totally-unknown-model"],
|
||||
["flux1dev"], # "Flux.1 D" normalises to "flux1d", not "flux1"
|
||||
["sd"],
|
||||
["ponyxl"],
|
||||
],
|
||||
)
|
||||
def test_returns_empty_when_not_exactly_sure(self, hints):
|
||||
assert resolve_base_model(hints, KNOWN) == ""
|
||||
|
||||
def test_returns_empty_without_a_vocabulary(self):
|
||||
assert resolve_base_model(["KREA_2"], []) == ""
|
||||
|
||||
def test_only_ever_returns_a_known_name(self):
|
||||
for name in ["KREA_2", "KREA_2_TURBO", "krea-2-turbo", "unknown"]:
|
||||
result = resolve_base_model([name], KNOWN)
|
||||
assert result == "" or result in KNOWN
|
||||
|
||||
def test_real_world_modelscope_hints(self):
|
||||
"""The hints ModelScope actually publishes for a Krea 2 LoRA."""
|
||||
assert (
|
||||
resolve_base_model(
|
||||
["KREA_2", "KREA_2_TURBO", "Krea-2-Turbo"], KNOWN
|
||||
)
|
||||
== "Krea 2"
|
||||
)
|
||||
@@ -7,9 +7,13 @@ each provider's model-card fetching and capability flags.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest import mock
|
||||
import pytest
|
||||
|
||||
from py.services.agent.agent_service import AgentService
|
||||
from py.services.model_sources import (
|
||||
ModelCardContext,
|
||||
ModelSourceCache,
|
||||
HuggingFaceSource,
|
||||
ModelScopeSource,
|
||||
TensorArtSource,
|
||||
@@ -332,6 +336,206 @@ class TestAssetBaseUrl:
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model card context (site extras kept outside the README)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _modelscope_detail_payload() -> dict:
|
||||
"""A trimmed-but-faithful ModelScope model-detail response.
|
||||
|
||||
Mirrors the shape of ``/api/v1/models/{id}`` for an AIGC LoRA repo whose
|
||||
README is auto-generated boilerplate, so the author summary and the
|
||||
per-file example images are only reachable through this API.
|
||||
"""
|
||||
|
||||
return {
|
||||
"Code": 200,
|
||||
"Data": {
|
||||
"Name": "Krea-2-LORA",
|
||||
"ChineseName": "krea脸模",
|
||||
"Description": "权重0.5-1.2。配合《风格滤镜》lora一起使用。",
|
||||
"BaseModel": ["krea/Krea-2-Turbo"],
|
||||
"License": "Apache License 2.0",
|
||||
"OfficialTags": [
|
||||
{"Tag": "photography", "ChineseName": "写实摄影"},
|
||||
{"Tag": "woman", "ChineseName": "女生"},
|
||||
{"Tag": "photography", "ChineseName": "重复项"},
|
||||
],
|
||||
"MuseInfo": {
|
||||
"versions": [
|
||||
{
|
||||
"stats": {"fileList": ["Krea-2-LORA_c1-st8000.safetensors"]},
|
||||
"modelVersion": {"showName": "c1-st8000", "triggerWords": '[""]'},
|
||||
"coverImages": [
|
||||
{"url": "https://resources.modelscope.cn/cover-images/a.png"}
|
||||
],
|
||||
},
|
||||
{
|
||||
"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"},
|
||||
],
|
||||
},
|
||||
]
|
||||
},
|
||||
"ModelInfos": {
|
||||
"safetensor": {
|
||||
"files": [
|
||||
{
|
||||
"name": "Krea-2-LORA_c1-st8000.safetensors",
|
||||
"sha256": "a" * 64,
|
||||
"size": 234680568,
|
||||
},
|
||||
{
|
||||
"name": "Krea-2-LORA_c1-st1000.safetensors",
|
||||
"sha256": "b" * 64,
|
||||
"size": 234680568,
|
||||
},
|
||||
]
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestFetchModelCardContext:
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_reads_description_tags_and_base_model(self, monkeypatch):
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
assert url == "https://modelscope.cn/api/v1/models/u/r"
|
||||
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 context.description == "权重0.5-1.2。配合《风格滤镜》lora一起使用。"
|
||||
assert context.base_model == "krea/Krea-2-Turbo"
|
||||
# OfficialTag values only, de-duplicated, order preserved.
|
||||
assert context.official_tags == ["photography", "woman"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_matches_example_images_by_filename(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"
|
||||
)
|
||||
|
||||
# Only the requested file's images, never a sibling checkpoint's.
|
||||
assert context.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
"https://resources.modelscope.cn/cover-images/c.png",
|
||||
]
|
||||
assert context.trigger_words == ["kreaface", "kreamodel"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_never_borrows_images_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.example_images == []
|
||||
assert context.trigger_words == []
|
||||
# The repo-wide fields are still returned.
|
||||
assert context.base_model == "krea/Krea-2-Turbo"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_single_version_repo_without_filename(self, monkeypatch):
|
||||
payload = _modelscope_detail_payload()
|
||||
versions = payload["Data"]["MuseInfo"]["versions"]
|
||||
payload["Data"]["MuseInfo"]["versions"] = versions[:1]
|
||||
|
||||
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.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/a.png"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_tolerates_failures_and_odd_payloads(self, monkeypatch):
|
||||
payloads = (None, {"Code": 500}, {"Data": "nope"}, {"Data": {}})
|
||||
for payload in payloads:
|
||||
|
||||
async def fake_fetch_json(url, _payload=payload, **_kwargs):
|
||||
return (0 if _payload is None else 200), _payload
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
|
||||
)
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "a.safetensors"
|
||||
)
|
||||
assert context.is_empty(), payload
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_reads_stats_from_json_encoded_fallback(self, monkeypatch):
|
||||
payload = {
|
||||
"Data": {
|
||||
"MuseInfo": {
|
||||
"versions": [
|
||||
{
|
||||
"modelVersion": {
|
||||
"showName": "v1",
|
||||
"stats": '{"fileList": ["model.safetensors"]}',
|
||||
"triggerWords": '["hi"]',
|
||||
},
|
||||
"coverImages": [{"url": "https://cdn.example/x.png"}],
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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", "model.safetensors"
|
||||
)
|
||||
|
||||
assert context.example_images == ["https://cdn.example/x.png"]
|
||||
assert context.trigger_words == ["hi"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_context_is_empty_for_other_sources(self):
|
||||
assert (await HuggingFaceSource().fetch_model_card_context("u/r")).is_empty()
|
||||
assert (await TensorArtSource().fetch_model_card_context("123")).is_empty()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Download support
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -504,3 +708,251 @@ class TestDownloadSourceRegistry:
|
||||
assert get_download_source("nope") is None
|
||||
assert get_download_source("modelscope").platform == "modelscope"
|
||||
assert get_download_source("huggingface").platform == "huggingface"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-run model-card cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestModelSourceCache:
|
||||
def test_starts_empty(self):
|
||||
cache = ModelSourceCache()
|
||||
assert cache.readmes == {}
|
||||
assert cache.provider == {}
|
||||
|
||||
|
||||
class TestCachedModelCardFetch:
|
||||
@pytest.mark.asyncio
|
||||
async def test_detail_payload_is_fetched_once_per_source_id(self, monkeypatch):
|
||||
"""A collection repo's files share one detail request, not one each."""
|
||||
calls: list[str] = []
|
||||
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
calls.append(url)
|
||||
return 200, _modelscope_detail_payload()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
|
||||
)
|
||||
|
||||
source = ModelScopeSource()
|
||||
cache = ModelSourceCache()
|
||||
filenames = [
|
||||
"Krea-2-LORA_c1-st1000.safetensors",
|
||||
"Krea-2-LORA_c1-st8000.safetensors",
|
||||
"Krea-2-LORA_c1-st1000.safetensors",
|
||||
]
|
||||
for name in filenames:
|
||||
await source.fetch_model_card_context("u/r", name, cache=cache)
|
||||
|
||||
assert calls == ["https://modelscope.cn/api/v1/models/u/r"]
|
||||
assert ("modelscope", "detail", "u/r") in cache.provider
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_file_selection_still_runs_for_each_file(self, monkeypatch):
|
||||
"""The cached payload must not leak one file's images to another."""
|
||||
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
return 200, _modelscope_detail_payload()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
|
||||
)
|
||||
|
||||
source = ModelScopeSource()
|
||||
cache = ModelSourceCache()
|
||||
|
||||
st1000 = await source.fetch_model_card_context(
|
||||
"u/r", "Krea-2-LORA_c1-st1000.safetensors", cache=cache
|
||||
)
|
||||
st8000 = await source.fetch_model_card_context(
|
||||
"u/r", "Krea-2-LORA_c1-st8000.safetensors", cache=cache
|
||||
)
|
||||
|
||||
assert st1000.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
"https://resources.modelscope.cn/cover-images/c.png",
|
||||
]
|
||||
assert st8000.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/a.png"
|
||||
]
|
||||
assert st8000.trigger_words == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failures_are_not_cached(self, monkeypatch):
|
||||
"""A transient error must be retried for the next file."""
|
||||
calls: list[str] = []
|
||||
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
calls.append(url)
|
||||
return 500, None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
|
||||
)
|
||||
|
||||
source = ModelScopeSource()
|
||||
cache = ModelSourceCache()
|
||||
for _ in range(2):
|
||||
context = await source.fetch_model_card_context("u/r", "a.safetensors", cache=cache)
|
||||
assert context.is_empty()
|
||||
|
||||
assert len(calls) == 2
|
||||
assert cache.provider == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_cache_keeps_the_uncached_behaviour(self, monkeypatch):
|
||||
calls: list[str] = []
|
||||
|
||||
async def fake_fetch_json(url, **_kwargs):
|
||||
calls.append(url)
|
||||
return 200, _modelscope_detail_payload()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.model_sources.modelscope.fetch_json", fake_fetch_json
|
||||
)
|
||||
|
||||
source = ModelScopeSource()
|
||||
for _ in range(2):
|
||||
await source.fetch_model_card_context("u/r", "a.safetensors")
|
||||
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
class TestHashBasedVersionMatching:
|
||||
"""A renamed file must still find its own example images."""
|
||||
|
||||
ST1000_HASH = "b" * 64
|
||||
ST8000_HASH = "a" * 64
|
||||
|
||||
@staticmethod
|
||||
def _patch(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
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_renamed_file_is_matched_by_sha256(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "krea脸模-st1000-我改的名字.safetensors", sha256=self.ST1000_HASH
|
||||
)
|
||||
|
||||
assert context.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
"https://resources.modelscope.cn/cover-images/c.png",
|
||||
]
|
||||
assert context.trigger_words == ["kreaface", "kreamodel"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_renamed_file_without_a_hash_finds_nothing(self, monkeypatch):
|
||||
"""Pins the behaviour the hash match exists to fix."""
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "krea脸模-st1000-我改的名字.safetensors"
|
||||
)
|
||||
|
||||
assert context.example_images == []
|
||||
# Repo-wide fields are unaffected by the miss.
|
||||
assert context.base_model == "krea/Krea-2-Turbo"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hash_wins_over_a_filename_that_matches_another_version(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""An inconsistent name/hash pair trusts the content hash."""
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "Krea-2-LORA_c1-st8000.safetensors", sha256=self.ST1000_HASH
|
||||
)
|
||||
|
||||
assert context.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
"https://resources.modelscope.cn/cover-images/c.png",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_hash_falls_back_to_the_filename(self, monkeypatch):
|
||||
"""A re-encoded file still matches by name rather than losing its images."""
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "Krea-2-LORA_c1-st1000.safetensors", sha256="f" * 64
|
||||
)
|
||||
|
||||
assert context.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
"https://resources.modelscope.cn/cover-images/c.png",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hash_is_case_insensitive(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "renamed.safetensors", sha256=self.ST1000_HASH.upper()
|
||||
)
|
||||
|
||||
assert len(context.example_images) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blank_hash_is_ignored(self, monkeypatch):
|
||||
self._patch(monkeypatch)
|
||||
|
||||
context = await ModelScopeSource().fetch_model_card_context(
|
||||
"u/r", "Krea-2-LORA_c1-st8000.safetensors", sha256=" "
|
||||
)
|
||||
|
||||
assert context.example_images == [
|
||||
"https://resources.modelscope.cn/cover-images/a.png"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_model_infos_degrades_to_filename_matching(self, monkeypatch):
|
||||
payload = _modelscope_detail_payload()
|
||||
del payload["Data"]["ModelInfos"]
|
||||
|
||||
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", "Krea-2-LORA_c1-st1000.safetensors", sha256=self.ST1000_HASH
|
||||
)
|
||||
|
||||
assert len(context.example_images) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_published_hashes_are_used_by_the_agent(self):
|
||||
"""The agent must pass the recorded hash, not just the filename."""
|
||||
service = AgentService()
|
||||
with (
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
|
||||
new=mock.AsyncMock(return_value="# card"),
|
||||
),
|
||||
mock.patch(
|
||||
"py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card_context",
|
||||
new=mock.AsyncMock(return_value=ModelCardContext()),
|
||||
) as mock_ctx,
|
||||
):
|
||||
await service._load_source_card(
|
||||
"/models/loras/renamed.safetensors",
|
||||
{
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/u/r",
|
||||
"sha256": "c" * 64,
|
||||
},
|
||||
)
|
||||
|
||||
assert mock_ctx.call_args.kwargs["sha256"] == "c" * 64
|
||||
|
||||
@@ -6,12 +6,14 @@ functions and verify the business logic (conditions, merges, dispatch).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.agent.post_processor import PostProcessor
|
||||
from py.services.model_sources import ModelCardContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -437,6 +439,45 @@ Content
|
||||
assert applied["metadata_source"] == "agent:enrich_hf_metadata"
|
||||
assert "llm_enriched_at" in applied
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_confidence_is_stored_under_a_persisted_key(self, processor):
|
||||
"""`llm_confidence` must not be underscore-prefixed.
|
||||
|
||||
Underscore-prefixed keys are dropped by `BaseModelMetadata`, which made
|
||||
`_llm_confidence` vanish on the next metadata write.
|
||||
"""
|
||||
llm = {**self.MIN_LLM_OUTPUT, "confidence": "medium"}
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview", return_value=False),
|
||||
mock.patch("py.metadata_ops.refresh_cache"),
|
||||
):
|
||||
await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=llm,
|
||||
metadata={},
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert applied["llm_confidence"] == "medium"
|
||||
assert "_llm_confidence" not in applied
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_confidence_absent_when_the_llm_reported_none(self, processor):
|
||||
llm = {**self.MIN_LLM_OUTPUT, "confidence": ""}
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview", return_value=False),
|
||||
mock.patch("py.metadata_ops.refresh_cache"),
|
||||
):
|
||||
await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=llm,
|
||||
metadata={},
|
||||
)
|
||||
assert "llm_confidence" not in mock_apply.call_args[0][1]
|
||||
|
||||
# -- preview download ------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -523,3 +564,560 @@ class TestMergeTags:
|
||||
result = PostProcessor._merge_tags(existing, new)
|
||||
# All tags are lowercased (matching TagUpdateService behaviour)
|
||||
assert result == ["anime", "flux", "lora"]
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# enrich_hf_metadata — site-provided card extras (ModelCardContext)
|
||||
# ======================================================================
|
||||
|
||||
|
||||
class TestSiteProvidedContext:
|
||||
"""ModelScope keeps the author summary, the curated tags and the per-file
|
||||
example images outside the README; these tests pin how they are applied.
|
||||
"""
|
||||
|
||||
MODELSCOPE_METADATA = {
|
||||
"from_civitai": False,
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/user/repo",
|
||||
}
|
||||
|
||||
LLM_OUTPUT = {
|
||||
"base_model": "",
|
||||
"trigger_words": [],
|
||||
"short_description": "",
|
||||
"tags": [],
|
||||
"recommended_width": 0,
|
||||
"recommended_height": 0,
|
||||
"preview_url": "",
|
||||
"confidence": "medium",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_example_images_become_gallery_and_preview(self, processor):
|
||||
"""A boilerplate README still yields images and a downloaded preview."""
|
||||
context = ModelCardContext(
|
||||
example_images=[
|
||||
"https://resources.modelscope.cn/cover-images/a.png",
|
||||
"https://resources.modelscope.cn/cover-images/b.png",
|
||||
]
|
||||
)
|
||||
boilerplate = "### 当前模型的贡献者未提供更加详细的模型介绍。\n"
|
||||
|
||||
with (
|
||||
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.metadata_ops.download_preview") as mock_dl,
|
||||
mock.patch("py.metadata_ops.refresh_cache"),
|
||||
):
|
||||
mock_dl.return_value = "/p.webp"
|
||||
result = await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
llm_output=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=boilerplate,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
applied = mock_apply.call_args[0][1]
|
||||
images = applied["civitai"]["images"]
|
||||
assert [img["url"] for img in images] == context.example_images
|
||||
assert images[0]["type"] == "image"
|
||||
# The first (per-file) site image is used as the preview.
|
||||
mock_dl.assert_awaited_once_with(
|
||||
"/p.safetensors", "https://resources.modelscope.cn/cover-images/a.png"
|
||||
)
|
||||
assert applied["preview_url"] == "/p.webp"
|
||||
assert result["preview_downloaded"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_example_images_work_without_any_readme(self, processor):
|
||||
"""The site images alone are enough — the README may be unreachable."""
|
||||
context = ModelCardContext(
|
||||
example_images=["https://resources.modelscope.cn/cover-images/a.png"]
|
||||
)
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content="",
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
images = mock_apply.call_args[0][1]["civitai"]["images"]
|
||||
assert [img["url"] for img in images] == context.example_images
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_description_precedes_readme_in_model_description(self, processor):
|
||||
context = ModelCardContext(description="权重0.5-1.2。配合滤镜lora一起使用。")
|
||||
readme = "# 模型介绍\n\n本模型依托魔搭社区完成训练。\n"
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=readme,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
description = mock_apply.call_args[0][1]["modelDescription"]
|
||||
assert description.startswith(f"<p>{context.description}</p>")
|
||||
assert "<h1>模型介绍</h1>" in description
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_description_is_html_escaped(self, processor):
|
||||
context = ModelCardContext(description="a < b & c")
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content="",
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
assert mock_apply.call_args[0][1]["modelDescription"] == "<p>a < b & c</p>"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_trigger_words_fill_in_when_llm_finds_none(self, processor):
|
||||
context = ModelCardContext(trigger_words=["kreaface", "kreamodel"])
|
||||
readme = "---\ninstance_prompt: yamlword\n---\nbody\n"
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=readme,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
# The per-file site value wins over the repo-wide YAML instance_prompt.
|
||||
assert mock_apply.call_args[0][1]["civitai"]["trainedWords"] == [
|
||||
"kreaface",
|
||||
"kreamodel",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_yaml_instance_prompt_still_used_when_site_has_none(self, processor):
|
||||
context = ModelCardContext(description="summary only")
|
||||
readme = "---\ninstance_prompt: yamlword\n---\nbody\n"
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=readme,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
assert mock_apply.call_args[0][1]["civitai"]["trainedWords"] == ["yamlword"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_images_are_skipped_for_a_model_with_no_external_source(
|
||||
self, processor
|
||||
):
|
||||
"""A CivitAI-only model must not pick up ModelScope images."""
|
||||
context = ModelCardContext(
|
||||
example_images=["https://resources.modelscope.cn/cover-images/a.png"]
|
||||
)
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata={"from_civitai": True},
|
||||
readme_content="",
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
assert "images" not in mock_apply.call_args[0][1].get("civitai", {})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_images_deduplicate_against_readme_images(self, processor):
|
||||
"""A URL present in both the site data and the README appears once."""
|
||||
shared = "https://modelscope.cn/models/user/repo/resolve/master/sample.png"
|
||||
context = ModelCardContext(example_images=[shared])
|
||||
readme = f"\n"
|
||||
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=readme,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
images = mock_apply.call_args[0][1]["civitai"]["images"]
|
||||
assert [img["url"] for img in images] == [shared]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_context_keeps_readme_only_behaviour(self, processor):
|
||||
"""An empty site context must not change existing HF behaviour."""
|
||||
readme = "---\nwidget:\n- text: a cat\n output:\n url: images/cat.png\n---\n"
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata={
|
||||
"from_civitai": False,
|
||||
"hf_url": "https://huggingface.co/user/repo",
|
||||
},
|
||||
readme_content=readme,
|
||||
source_context=ModelCardContext(),
|
||||
)
|
||||
images = mock_apply.call_args[0][1]["civitai"]["images"]
|
||||
assert [img["url"] for img in images] == [
|
||||
"https://huggingface.co/user/repo/resolve/main/images/cat.png"
|
||||
]
|
||||
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# enrich_hf_metadata — deterministic fallbacks used when the LLM is skipped
|
||||
# ======================================================================
|
||||
|
||||
|
||||
class TestDeterministicFallbacks:
|
||||
"""With the LLM skipped, these fields must still be produced from the API."""
|
||||
|
||||
MODELSCOPE_METADATA = {
|
||||
"from_civitai": False,
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/user/repo",
|
||||
}
|
||||
|
||||
EMPTY_LLM = {
|
||||
"base_model": "",
|
||||
"trigger_words": [],
|
||||
"short_description": "",
|
||||
"tags": [],
|
||||
"recommended_width": 0,
|
||||
"recommended_height": 0,
|
||||
"preview_url": "",
|
||||
"notes": "",
|
||||
"usage_tips": "{}",
|
||||
"confidence": "",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolved_base_model_used_when_llm_gave_none(self, processor):
|
||||
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=self.EMPTY_LLM,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=ModelCardContext(base_model="krea/Krea-2-Turbo"),
|
||||
resolved_base_model="Krea 2",
|
||||
)
|
||||
assert mock_apply.call_args[0][1]["base_model"] == "Krea 2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_base_model_still_wins_over_the_resolver(self, processor):
|
||||
llm = {**self.EMPTY_LLM, "base_model": "Flux.1 D"}
|
||||
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,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
resolved_base_model="Krea 2",
|
||||
)
|
||||
assert mock_apply.call_args[0][1]["base_model"] == "Flux.1 D"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_site_description_fills_civitai_description(self, processor):
|
||||
context = ModelCardContext(description="一个 Krea 2 人像 LoRA。")
|
||||
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=self.EMPTY_LLM,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=context,
|
||||
)
|
||||
assert (
|
||||
mock_apply.call_args[0][1]["civitai"]["description"]
|
||||
== "一个 Krea 2 人像 LoRA。"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_short_description_wins_over_site_description(self, processor):
|
||||
llm = {**self.EMPTY_LLM, "short_description": "from the LLM"}
|
||||
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,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=ModelCardContext(description="from the site"),
|
||||
)
|
||||
assert mock_apply.call_args[0][1]["civitai"]["description"] == "from the LLM"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_official_tags_are_applied_without_the_llm(self, processor):
|
||||
context = ModelCardContext(
|
||||
official_tags=["photography", "character-enhancement", "woman"]
|
||||
)
|
||||
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=self.EMPTY_LLM,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=context,
|
||||
)
|
||||
assert mock_apply.call_args[0][1]["tags"] == [
|
||||
"photography",
|
||||
"character-enhancement",
|
||||
"woman",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_official_tags_are_kept_alongside_llm_tags(self, processor):
|
||||
context = ModelCardContext(official_tags=["photography", "woman"])
|
||||
llm = {**self.EMPTY_LLM, "tags": ["portrait", "photography"]}
|
||||
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,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=context,
|
||||
)
|
||||
# Site tags first, then the LLM's extra ones, no duplicates.
|
||||
assert mock_apply.call_args[0][1]["tags"] == [
|
||||
"photography",
|
||||
"woman",
|
||||
"portrait",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_tips_recovered_from_the_author_summary(self, processor):
|
||||
context = ModelCardContext(
|
||||
description="权重0.5-1.2。2个一起时,权重建议都用1.0-1.1。"
|
||||
)
|
||||
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=self.EMPTY_LLM,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=context,
|
||||
)
|
||||
tips = json.loads(mock_apply.call_args[0][1]["usage_tips"])
|
||||
assert tips == {
|
||||
"strength_min": 0.5,
|
||||
"strength_max": 1.2,
|
||||
"strength_range": "0.5-1.2",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_usage_tips_win_over_the_regex(self, processor):
|
||||
llm = {**self.EMPTY_LLM, "usage_tips": '{"strength": 0.9}'}
|
||||
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,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=ModelCardContext(description="权重0.5-1.2"),
|
||||
)
|
||||
assert mock_apply.call_args[0][1]["usage_tips"] == '{"strength": 0.9}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notes_are_not_rewritten_when_the_llm_is_skipped(self, processor):
|
||||
"""Notes are LLM-only; skipping must not clobber or duplicate them."""
|
||||
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=self.EMPTY_LLM,
|
||||
metadata={**self.MODELSCOPE_METADATA, "notes": "existing notes"},
|
||||
source_context=ModelCardContext(description="权重0.5-1.2"),
|
||||
)
|
||||
assert "notes" not in mock_apply.call_args[0][1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_site_data_leaves_llm_only_fields_untouched(self, processor):
|
||||
"""An empty context must behave exactly like the pre-existing pipeline."""
|
||||
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=self.EMPTY_LLM,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
source_context=ModelCardContext(),
|
||||
resolved_base_model="",
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert "base_model" not in applied
|
||||
assert "tags" not in applied
|
||||
assert "notes" not in applied
|
||||
assert "usage_tips" not in applied
|
||||
|
||||
|
||||
class TestPlaceholderCardDescription:
|
||||
"""A site-generated placeholder card must not become the description."""
|
||||
|
||||
MODELSCOPE_METADATA = {
|
||||
"from_civitai": False,
|
||||
"source_platform": "modelscope",
|
||||
"source_url": "https://modelscope.cn/models/user/repo",
|
||||
}
|
||||
|
||||
PLACEHOLDER_README = """---
|
||||
base_model: krea/Krea-2-Turbo
|
||||
---
|
||||
### 当前模型的贡献者未提供更加详细的模型介绍。模型文件和权重,可浏览“模型文件”页面获取。
|
||||
#### 您可以通过如下git clone命令,或者ModelScope SDK来下载模型
|
||||
|
||||
SDK下载
|
||||
```bash
|
||||
pip install modelscope
|
||||
```
|
||||
|
||||
<p style="color: lightgrey;">如果您是本模型的贡献者,我们邀请您根据文档及时完善模型卡片内容。</p>
|
||||
"""
|
||||
|
||||
LLM_OUTPUT = {
|
||||
"base_model": "Krea 2",
|
||||
"trigger_words": [],
|
||||
"short_description": "一个 Krea 2 人像 LoRA。",
|
||||
"tags": [],
|
||||
"recommended_width": 0,
|
||||
"recommended_height": 0,
|
||||
"preview_url": "",
|
||||
"notes": "",
|
||||
"usage_tips": "{}",
|
||||
"confidence": "medium",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_description_holds_only_the_author_summary(self, processor):
|
||||
context = ModelCardContext(description="权重0.5-1.2。")
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=self.PLACEHOLDER_README,
|
||||
source_context=context,
|
||||
)
|
||||
|
||||
description = mock_apply.call_args[0][1]["modelDescription"]
|
||||
assert description == "<p>权重0.5-1.2。</p>"
|
||||
assert "pip install modelscope" not in description
|
||||
assert "git clone" not in description
|
||||
assert "贡献者" not in description
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_placeholder_card_alone_writes_no_description(self, processor):
|
||||
"""Without an author summary there is nothing worth storing."""
|
||||
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=self.LLM_OUTPUT,
|
||||
metadata=dict(self.MODELSCOPE_METADATA),
|
||||
readme_content=self.PLACEHOLDER_README,
|
||||
source_context=ModelCardContext(),
|
||||
)
|
||||
|
||||
assert "modelDescription" not in mock_apply.call_args[0][1]
|
||||
|
||||
@@ -212,3 +212,45 @@ async def test_self_healed_sidecar_is_parseable(tmp_path) -> None:
|
||||
assert metadata.file_name == "MyModel"
|
||||
assert metadata.model_name == "My Model"
|
||||
assert metadata.sha256 == "abc123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provenance_fields_survive_a_load_save_round_trip(tmp_path) -> None:
|
||||
"""Non-underscore extras must persist; underscore keys are ephemeral.
|
||||
|
||||
Regression test for `llm_confidence`: it was written as `_llm_confidence`,
|
||||
which `BaseModelMetadata.from_dict()` drops (and `to_dict()` strips), so the
|
||||
value was erased by the next metadata write and was invisible to
|
||||
`read_metadata()`. The enrichment evaluation harness depends on it.
|
||||
"""
|
||||
model_path = tmp_path / "Model.safetensors"
|
||||
model_path.write_bytes(b"fake model data")
|
||||
metadata_path = tmp_path / "Model.metadata.json"
|
||||
|
||||
payload = {
|
||||
"file_path": str(model_path),
|
||||
"file_name": "Model",
|
||||
"model_name": "Model",
|
||||
"sha256": "deadbeef",
|
||||
"base_model": "Krea 2",
|
||||
"preview_url": "",
|
||||
"metadata_source": "agent:enrich_hf_metadata",
|
||||
"llm_enriched_at": "2026-01-01T00:00:00+00:00",
|
||||
"llm_confidence": "medium",
|
||||
"_llm_confidence": "medium",
|
||||
}
|
||||
assert await MetadataManager.save_metadata(str(model_path), payload) is True
|
||||
|
||||
# A read must surface the supported key...
|
||||
loaded = await MetadataManager.load_metadata_payload(str(model_path))
|
||||
assert loaded["llm_confidence"] == "medium"
|
||||
assert loaded["metadata_source"] == "agent:enrich_hf_metadata"
|
||||
|
||||
# ...but the underscore alias is intentionally not persisted.
|
||||
assert "_llm_confidence" not in loaded
|
||||
|
||||
# Re-saving what we read must not lose the confidence.
|
||||
assert await MetadataManager.save_metadata(str(model_path), loaded) is True
|
||||
saved = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||
assert saved["llm_confidence"] == "medium"
|
||||
assert saved["llm_enriched_at"] == "2026-01-01T00:00:00+00:00"
|
||||
|
||||
Reference in New Issue
Block a user