mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 09:21:16 -03:00
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:
@@ -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") == ""
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user