feat(agent): extract HF widget gallery images into civitai.images with recommended dimensions

- Add extract_gallery_images() to parse YAML widget entries from README
  frontmatter, convert relative image URLs to absolute HF URLs, and
  build civitai.images-compatible entries with prompt metadata
- LLM now extracts recommended_width/recommended_height from README
  (e.g. "Best Dimensions"), used as gallery image dimensions
- extract_gallery_images() accepts default_width/height parameters,
  falling back to 512x512 when LLM provides no recommendation
- Frontend ShowcaseView.js: defensive NaN guard for 0 width/height
- post_processor: consistently merge civitai updates across triggers,
  description, and gallery blocks with distinct variable names
- SKILL.md: add recommended_width/recommended_height to output schema
- 62 tests pass, including gallery extraction and dimension tests
This commit is contained in:
Will Miao
2026-07-03 07:07:19 +08:00
parent 88349bf944
commit ee8250c26c
6 changed files with 314 additions and 9 deletions

View File

@@ -469,3 +469,100 @@ class TestConvertReadmeToHtml:
result = convert_readme_to_html(md)
assert "<pre>" not in result
assert "<h2>heading after spacing</h2>" in result
# ======================================================================
# extract_gallery_images — YAML widget → civitai.images
# ======================================================================
class TestExtractGalleryImages:
_REPO = "prithivMLmods/Flux-Long-Toon-LoRA"
_README = """---
tags:
- lora
widget:
- text: "a cat"
output:
url: images/cat.png
- text: >-
multi line
prompt here
output:
url: images/dog.png
base_model: flux
---
# Content after frontmatter
"""
def test_extracts_widget_images(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
images = extract_gallery_images(self._README, self._REPO)
assert len(images) == 2
assert images[0]["url"] == (
"https://huggingface.co/prithivMLmods/Flux-Long-Toon-LoRA"
"/resolve/main/images/cat.png"
)
assert images[0]["meta"]["prompt"] == "a cat"
assert images[0]["type"] == "image"
assert images[0]["hasMeta"] is True
assert images[0]["hasPositivePrompt"] is True
assert images[1]["url"] == (
"https://huggingface.co/prithivMLmods/Flux-Long-Toon-LoRA"
"/resolve/main/images/dog.png"
)
assert images[1]["meta"]["prompt"] == "multi line prompt here"
def test_default_dimensions_used(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
images = extract_gallery_images(self._README, self._REPO)
assert images[0]["width"] == 512
assert images[0]["height"] == 512
def test_custom_dimensions_applied(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
images = extract_gallery_images(
self._README, self._REPO,
default_width=768, default_height=1024,
)
assert images[0]["width"] == 768
assert images[0]["height"] == 1024
def test_empty_readme_returns_empty(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
assert extract_gallery_images("", self._REPO) == []
assert extract_gallery_images("no frontmatter here", self._REPO) == []
def test_empty_repo_returns_empty(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
assert extract_gallery_images(self._README, "") == []
def test_no_widget_returns_empty(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_gallery_images
md = "---\ntags:\n - lora\n---\n\nContent"
assert extract_gallery_images(md, self._REPO) == []
def test_extract_repo_from_hf_url(self):
from py.services.agent.skills.enrich_hf_metadata.md_to_html import \
extract_repo_from_hf_url
assert extract_repo_from_hf_url(
"https://huggingface.co/prithivMLmods/Flux-Long-Toon-LoRA"
) == "prithivMLmods/Flux-Long-Toon-LoRA"
assert extract_repo_from_hf_url("") == ""
assert extract_repo_from_hf_url("not a url") == ""

View File

@@ -69,6 +69,8 @@ class TestEnrichHfMetadata:
"trigger_words": [],
"short_description": "",
"tags": [],
"recommended_width": 0,
"recommended_height": 0,
"preview_url": "",
"confidence": "low",
}
@@ -244,6 +246,64 @@ class TestEnrichHfMetadata:
applied = mock_apply.call_args[0][1]
assert "modelDescription" not in applied
# -- gallery images → civitai.images ---------------------------------
@pytest.mark.asyncio
async def test_gallery_images_extracted_from_readme(self, processor):
"""Widget entries in README → civitai.images."""
readme = """---
widget:
- text: "a cat"
output:
url: images/cat.png
---
Content
"""
with (
mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply,
mock.patch("py.agent_cli.download_preview", return_value=None),
mock.patch("py.agent_cli.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": False,
"hf_url": "https://huggingface.co/user/repo",
},
readme_content=readme,
)
applied = mock_apply.call_args[0][1]
images = applied.get("civitai", {}).get("images", [])
assert len(images) == 1
assert images[0]["url"] == (
"https://huggingface.co/user/repo/resolve/main/images/cat.png"
)
assert images[0]["meta"]["prompt"] == "a cat"
@pytest.mark.asyncio
async def test_gallery_images_skipped_for_civitai_model(self, processor):
"""Gallery images NOT extracted for CivitAI models."""
with (
mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply,
mock.patch("py.agent_cli.download_preview", return_value=None),
mock.patch("py.agent_cli.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": True,
"hf_url": "https://huggingface.co/user/repo",
},
readme_content="---\nwidget:\n- text: a\n output:\n url: x.png\n---\n",
)
applied = mock_apply.call_args[0][1]
civitai = applied.get("civitai", {})
assert "images" not in civitai
# -- tags ------------------------------------------------------------
@pytest.mark.asyncio