fix(civitai): strip comfy metadata from images

This commit is contained in:
pixelpaws
2025-09-26 08:55:46 +08:00
parent 141f83065f
commit e7316b3389
2 changed files with 37 additions and 4 deletions

View File

@@ -33,6 +33,24 @@ class CivitaiClient:
self.base_url = "https://civitai.com/api/v1" self.base_url = "https://civitai.com/api/v1"
@staticmethod
def _remove_comfy_metadata(model_version: Optional[Dict]) -> None:
"""Remove Comfy-specific metadata from model version images."""
if not isinstance(model_version, dict):
return
images = model_version.get("images")
if not isinstance(images, list):
return
for image in images:
if not isinstance(image, dict):
continue
meta = image.get("meta")
if isinstance(meta, dict) and "comfy" in meta:
meta.pop("comfy", None)
async def download_file(self, url: str, save_dir: str, default_filename: str, progress_callback=None) -> Tuple[bool, str]: async def download_file(self, url: str, save_dir: str, default_filename: str, progress_callback=None) -> Tuple[bool, str]:
"""Download file with resumable downloads and retry mechanism """Download file with resumable downloads and retry mechanism
@@ -85,6 +103,7 @@ class CivitaiClient:
# Add creator from model data # Add creator from model data
result['creator'] = data.get("creator") result['creator'] = data.get("creator")
self._remove_comfy_metadata(result)
return result, None return result, None
# Handle specific error cases # Handle specific error cases
@@ -178,6 +197,7 @@ class CivitaiClient:
version['model']['tags'] = model_data.get("tags", []) version['model']['tags'] = model_data.get("tags", [])
version['creator'] = model_data.get("creator") version['creator'] = model_data.get("creator")
self._remove_comfy_metadata(version)
return version return version
# Case 2: model_id is provided (with or without version_id) # Case 2: model_id is provided (with or without version_id)
@@ -260,6 +280,7 @@ class CivitaiClient:
# Add creator from model data # Add creator from model data
version['creator'] = data.get("creator") version['creator'] = data.get("creator")
self._remove_comfy_metadata(version)
return version return version
# Case 3: Neither model_id nor version_id provided # Case 3: Neither model_id nor version_id provided
@@ -295,6 +316,7 @@ class CivitaiClient:
if success: if success:
logger.debug(f"Successfully fetched model version info for: {version_id}") logger.debug(f"Successfully fetched model version info for: {version_id}")
self._remove_comfy_metadata(result)
return result, None return result, None
# Handle specific error cases # Handle specific error cases

View File

@@ -63,6 +63,10 @@ async def test_get_model_by_hash_enriches_metadata(monkeypatch, downloader):
"modelId": 123, "modelId": 123,
"model": {"description": "", "tags": []}, "model": {"description": "", "tags": []},
"creator": {}, "creator": {},
"images": [
{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}},
{"meta": "not-a-dict"},
],
} }
model_payload = {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}} model_payload = {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}}
@@ -83,6 +87,8 @@ async def test_get_model_by_hash_enriches_metadata(monkeypatch, downloader):
assert result["model"]["description"] == "desc" assert result["model"]["description"] == "desc"
assert result["model"]["tags"] == ["tag"] assert result["model"]["tags"] == ["tag"]
assert result["creator"] == {"username": "user"} assert result["creator"] == {"username": "user"}
assert "comfy" not in result["images"][0]["meta"]
assert result["images"][0]["meta"]["other"] == "keep"
async def test_get_model_by_hash_handles_not_found(monkeypatch, downloader): async def test_get_model_by_hash_handles_not_found(monkeypatch, downloader):
@@ -144,6 +150,7 @@ async def test_get_model_version_by_version_id(monkeypatch, downloader):
"modelId": 321, "modelId": 321,
"model": {"description": ""}, "model": {"description": ""},
"files": [], "files": [],
"images": [{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}}],
} }
if url.endswith("/models/321"): if url.endswith("/models/321"):
return True, {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}} return True, {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}}
@@ -158,6 +165,8 @@ async def test_get_model_version_by_version_id(monkeypatch, downloader):
assert result["model"]["description"] == "desc" assert result["model"]["description"] == "desc"
assert result["model"]["tags"] == ["tag"] assert result["model"]["tags"] == ["tag"]
assert result["creator"] == {"username": "user"} assert result["creator"] == {"username": "user"}
assert "comfy" not in result["images"][0]["meta"]
assert result["images"][0]["meta"]["other"] == "keep"
async def test_get_model_version_requires_identifier(monkeypatch, downloader): async def test_get_model_version_requires_identifier(monkeypatch, downloader):
@@ -181,7 +190,7 @@ async def test_get_model_version_info_handles_not_found(monkeypatch, downloader)
async def test_get_model_version_info_success(monkeypatch, downloader): async def test_get_model_version_info_success(monkeypatch, downloader):
expected = {"id": 55} expected = {"id": 55, "images": [{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}}]}
async def fake_make_request(method, url, use_auth=True): async def fake_make_request(method, url, use_auth=True):
return True, expected return True, expected
@@ -194,6 +203,8 @@ async def test_get_model_version_info_success(monkeypatch, downloader):
assert result == expected assert result == expected
assert error is None assert error is None
assert "comfy" not in result["images"][0]["meta"]
assert result["images"][0]["meta"]["other"] == "keep"
async def test_get_image_info_returns_first_item(monkeypatch, downloader): async def test_get_image_info_returns_first_item(monkeypatch, downloader):