mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 09:21:16 -03:00
feat(agent): improve enrich_hf_metadata skill with priority_tags, preview_url fix, civitai.trainedWords
- Add identify_model_type() helper to determine lora/checkpoint/embedding - Pass priority_tags from user settings to LLM prompt for tag relevance - SKILL.md: instruct LLM to exclude technical/generic HF tags, cross-reference against priority_tags; forbid ['None'] placeholder for trigger words - post_processor: fix preview_url not updated after download (now writes local .webp path to metadata); write trigger words to civitai.trainedWords instead of top-level; sanitize ['None']/'null'/'n/a' placeholder values to [] - download_preview() now returns str | None (local path) instead of bool - Update tests for new return type and nested civitai.trainedWords structure
This commit is contained in:
@@ -227,11 +227,11 @@ class TestApplyMetadataUpdates:
|
||||
class TestDownloadPreview:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_url_returns_false(self, tmp_path):
|
||||
async def test_empty_url_returns_none(self, tmp_path):
|
||||
mp = tmp_path / "m.safetensors"
|
||||
mp.write_bytes(b"fake")
|
||||
assert await download_preview(str(mp), "") is False
|
||||
assert await download_preview(str(mp), " ") is False
|
||||
assert await download_preview(str(mp), "") is None
|
||||
assert await download_preview(str(mp), " ") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_download_and_optimise(self, tmp_path):
|
||||
@@ -246,12 +246,12 @@ class TestDownloadPreview:
|
||||
get_dl.return_value = dl
|
||||
exif.optimize_image.return_value = (b"optimized_webp", {})
|
||||
result = await download_preview(str(mp), "https://ex.com/i.png")
|
||||
assert result is True
|
||||
assert result == str(tmp_path / "t.webp")
|
||||
assert (tmp_path / "t.webp").exists()
|
||||
assert (tmp_path / "t.webp").read_bytes() == b"optimized_webp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_failure_returns_false(self, tmp_path):
|
||||
async def test_download_failure_returns_none(self, tmp_path):
|
||||
mp = tmp_path / "t.safetensors"
|
||||
mp.write_bytes(b"fake")
|
||||
with mock.patch("py.services.downloader.get_downloader") as get_dl:
|
||||
@@ -260,7 +260,7 @@ class TestDownloadPreview:
|
||||
dl.download_file = mock.AsyncMock(return_value=(False, None))
|
||||
get_dl.return_value = dl
|
||||
result = await download_preview(str(mp), "https://ex.com/i.png")
|
||||
assert result is False
|
||||
assert result is None
|
||||
assert not (tmp_path / "t.webp").exists()
|
||||
|
||||
|
||||
|
||||
@@ -44,7 +44,7 @@ class TestProcessDispatch:
|
||||
mock.patch("py.agent_cli.refresh_cache") as mock_ref,
|
||||
):
|
||||
mock_apply.return_value = ["metadata_source"]
|
||||
mock_dl.return_value = False
|
||||
mock_dl.return_value = None
|
||||
|
||||
result = await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
@@ -155,7 +155,7 @@ class TestEnrichHfMetadata:
|
||||
llm = {**self.MIN_LLM_OUTPUT, "trigger_words": ["trigger1", "trigger2"]}
|
||||
with (
|
||||
mock.patch("py.agent_cli.apply_metadata_updates") as mock_apply,
|
||||
mock.patch("py.agent_cli.download_preview", return_value=False),
|
||||
mock.patch("py.agent_cli.download_preview", return_value=None),
|
||||
mock.patch("py.agent_cli.refresh_cache"),
|
||||
):
|
||||
await processor.process(
|
||||
@@ -165,7 +165,7 @@ class TestEnrichHfMetadata:
|
||||
metadata={"trainedWords": []},
|
||||
)
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert applied["trainedWords"] == ["trigger1", "trigger2"]
|
||||
assert applied["civitai"]["trainedWords"] == ["trigger1", "trigger2"]
|
||||
|
||||
# -- description -----------------------------------------------------
|
||||
|
||||
@@ -237,7 +237,7 @@ class TestEnrichHfMetadata:
|
||||
mock.patch("py.agent_cli.download_preview") as mock_dl,
|
||||
mock.patch("py.agent_cli.refresh_cache"),
|
||||
):
|
||||
mock_dl.return_value = True
|
||||
mock_dl.return_value = "/p.webp"
|
||||
result = await processor.process(
|
||||
skill_name="enrich_hf_metadata",
|
||||
model_path="/p.safetensors",
|
||||
@@ -246,6 +246,8 @@ class TestEnrichHfMetadata:
|
||||
)
|
||||
assert result["preview_downloaded"] is True
|
||||
mock_dl.assert_awaited_once_with("/p.safetensors", "https://ex.com/img.png")
|
||||
applied = mock_apply.call_args[0][1]
|
||||
assert applied["preview_url"] == "/p.webp"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_skipped_when_exists(self, processor):
|
||||
|
||||
Reference in New Issue
Block a user