diff --git a/py/services/agent/post_processor.py b/py/services/agent/post_processor.py index 16764cc4..d8f12d6e 100644 --- a/py/services/agent/post_processor.py +++ b/py/services/agent/post_processor.py @@ -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"

{html.escape(site_description)}

") + 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 @@ -217,16 +278,22 @@ class PostProcessor: if 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 +327,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 +395,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. diff --git a/py/services/agent/skills/enrich_hf_metadata/prompt.md b/py/services/agent/skills/enrich_hf_metadata/prompt.md index f02abf82..b7dd45c8 100644 --- a/py/services/agent/skills/enrich_hf_metadata/prompt.md +++ b/py/services/agent/skills/enrich_hf_metadata/prompt.md @@ -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, no spaces, no hyphens** (use single words like `"photorealistic"`, `"anime"`, `"character"`). This rule applies to Latin-script tags; when the model's own tags are in another script (e.g. Chinese), keep them verbatim instead of dropping or translating them. + +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 `![alt](url)` 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 { diff --git a/py/services/model_sources/__init__.py b/py/services/model_sources/__init__.py index 925143f4..290b7ecc 100644 --- a/py/services/model_sources/__init__.py +++ b/py/services/model_sources/__init__.py @@ -11,6 +11,7 @@ from __future__ import annotations from .base import ( GROUP_PREFIXES, HTTP_TIMEOUT, + ModelCardContext, ModelSource, ModelSourceError, SourceRef, @@ -45,6 +46,7 @@ __all__ = [ "GROUP_PREFIXES", "HTTP_TIMEOUT", "LEGACY_HF_URL_FIELD", + "ModelCardContext", "ModelSource", "ModelSourceError", "HuggingFaceSource", diff --git a/py/services/model_sources/base.py b/py/services/model_sources/base.py index a0ecc824..b0b13a45 100644 --- a/py/services/model_sources/base.py +++ b/py/services/model_sources/base.py @@ -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,7 +24,7 @@ from __future__ import annotations import logging import os import re -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, 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. @@ -230,6 +282,22 @@ class ModelSource: return "" + async def fetch_model_card_context( + self, source_id: str, filename: str = "" + ) -> ModelCardContext: + """Return the card extras the site keeps outside the README. + + *filename* is the model file's basename (no directory) and selects + the right entry when a repository holds several models. 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,6 +369,7 @@ def filter_weight_files(entries: Iterable[tuple[str, int]]) -> list[dict[str, An __all__ = [ "GROUP_PREFIXES", "HTTP_TIMEOUT", + "ModelCardContext", "ModelSource", "ModelSourceError", "SourceRef", diff --git a/py/services/model_sources/modelscope.py b/py/services/model_sources/modelscope.py index c1acf198..d0324d6c 100644 --- a/py/services/model_sources/modelscope.py +++ b/py/services/model_sources/modelscope.py @@ -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. @@ -22,10 +28,14 @@ valid; the CDN URL must never be cached. from __future__ import annotations +import json import logging +import os import re +from typing import Any from .base import ( + ModelCardContext, ModelSource, ModelSourceError, fetch_json, @@ -95,6 +105,47 @@ class ModelScopeSource(ModelSource): return text return "" + async def fetch_model_card_context( + self, source_id: str, filename: str = "" + ) -> 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. + + Example images are matched to *filename* through each version's + ``stats.fileList``, which means the images returned belong to the + exact ``.safetensors`` being enriched — essential for collection + repositories, where every checkpoint has its own sample image. + """ + + 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 ModelCardContext() + data = payload.get("Data") + if not isinstance(data, dict): + return ModelCardContext() + + 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) + if versions: + context.example_images = _cover_image_urls(versions) + context.trigger_words = _version_trigger_words(versions) + return context + async def list_files( self, source_id: str, revision: str = "" ) -> list[dict]: @@ -142,3 +193,213 @@ 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 _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 _matching_versions(muse_info: Any, filename: str) -> list[dict[str, Any]]: + """Return the ``versions`` entries that publish *filename*. + + Matching is by exact basename first, then by the version's ``showName`` + appearing in the file stem (which absorbs the naming drift ModelScope + sometimes applies to uploaded weights). All matches are returned so a + file re-published across several versions contributes all of its + example images. With no *filename* 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 [] + + 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 diff --git a/tests/services/test_model_sources.py b/tests/services/test_model_sources.py index 0bf2ea2b..37cd7f9e 100644 --- a/tests/services/test_model_sources.py +++ b/tests/services/test_model_sources.py @@ -332,6 +332,190 @@ 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"}, + ], + }, + ] + }, + }, + } + + +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 # --------------------------------------------------------------------------- diff --git a/tests/services/test_post_processor.py b/tests/services/test_post_processor.py index c69aafd5..567aeb06 100644 --- a/tests/services/test_post_processor.py +++ b/tests/services/test_post_processor.py @@ -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 @@ -523,3 +525,481 @@ 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"

{context.description}

") + assert "

模型介绍

" 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"] == "

a < b & c

" + + @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"![alt]({shared})\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