diff --git a/py/recipes/parsers/civitai_image.py b/py/recipes/parsers/civitai_image.py index 87910c01..4864f9c8 100644 --- a/py/recipes/parsers/civitai_image.py +++ b/py/recipes/parsers/civitai_image.py @@ -42,6 +42,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser): "height", "Model", "Model hash", + "modelVersionIds", ) return any(key in payload for key in civitai_image_fields) @@ -429,6 +430,65 @@ class CivitaiApiMetadataParser(RecipeMetadataParser): result["loras"].append(lora_entry) + # Process modelVersionIds from Civitai image API + # These are model version IDs returned at root level when meta doesn't contain resources + if "modelVersionIds" in metadata and isinstance( + metadata["modelVersionIds"], list + ): + for version_id in metadata["modelVersionIds"]: + version_id_str = str(version_id) + + # Skip if we've already added this LoRA by version ID + if version_id_str in added_loras: + continue + + # Initialize lora entry with version ID + lora_entry = { + "id": version_id, + "modelId": 0, + "name": "Unknown LoRA", + "version": "", + "type": "lora", + "weight": 1.0, + "existsLocally": False, + "thumbnailUrl": "/loras_static/images/no-preview.png", + "baseModel": "", + "size": 0, + "downloadUrl": "", + "isDeleted": False, + } + + # Fetch model info from Civitai + if metadata_provider and version_id_str: + try: + civitai_info = ( + await metadata_provider.get_model_version_info( + version_id_str + ) + ) + + populated_entry = await self.populate_lora_from_civitai( + lora_entry, + civitai_info, + recipe_scanner, + base_model_counts, + ) + + if populated_entry is None: + continue # Skip invalid LoRA types + + lora_entry = populated_entry + except Exception as e: + logger.error( + f"Error fetching Civitai info for model version {version_id}: {e}" + ) + + # Track this LoRA for deduplication + if version_id_str: + added_loras[version_id_str] = len(result["loras"]) + + result["loras"].append(lora_entry) + # If we found LoRA hashes in the metadata but haven't already # populated entries for them, fall back to creating LoRAs from # the hashes section. Some Civitai image responses only include diff --git a/py/services/recipes/analysis_service.py b/py/services/recipes/analysis_service.py index 25932954..c2f4577a 100644 --- a/py/services/recipes/analysis_service.py +++ b/py/services/recipes/analysis_service.py @@ -143,6 +143,12 @@ class RecipeAnalysisService: ): metadata = metadata["meta"] + # Include modelVersionIds from root level if available + # Civitai API returns modelVersionIds at root level, not in meta + model_version_ids = image_info.get("modelVersionIds") + if model_version_ids and isinstance(metadata, dict): + metadata["modelVersionIds"] = model_version_ids + # Validate that metadata contains meaningful recipe fields # If not, treat as None to trigger EXIF extraction from downloaded image if isinstance(metadata, dict) and not self._has_recipe_fields(metadata): diff --git a/tests/services/test_civitai_image_parser.py b/tests/services/test_civitai_image_parser.py index f42782d8..6e8f5486 100644 --- a/tests/services/test_civitai_image_parser.py +++ b/tests/services/test_civitai_image_parser.py @@ -184,7 +184,10 @@ async def test_parse_metadata_populates_checkpoint_and_rewrites_thumbnails(monke assert result["model"] is not None assert result["model"]["name"] == "Checkpoint Example" assert result["model"]["type"] == "checkpoint" - assert result["model"]["thumbnailUrl"] == "https://image.civitai.com/checkpoints/width=450,optimized=true" + assert ( + result["model"]["thumbnailUrl"] + == "https://image.civitai.com/checkpoints/width=450,optimized=true" + ) assert result["model"]["modelId"] == 111 assert result["model"]["size"] == 1024 * 1024 assert result["model"]["hash"] == "ffaa0011" @@ -192,5 +195,106 @@ async def test_parse_metadata_populates_checkpoint_and_rewrites_thumbnails(monke assert result["loras"] assert result["loras"][0]["name"] == "Example Lora Model" - assert result["loras"][0]["thumbnailUrl"] == "https://image.civitai.com/loras/width=450,optimized=true" + assert ( + result["loras"][0]["thumbnailUrl"] + == "https://image.civitai.com/loras/width=450,optimized=true" + ) assert result["loras"][0]["hash"] == "abc123" + + +@pytest.mark.asyncio +async def test_parse_metadata_handles_modelVersionIds(monkeypatch): + """Test that modelVersionIds from Civitai image API are properly processed.""" + lora_info_1 = { + "id": 2398829, + "modelId": 123456, + "model": {"name": "Dance LoRA 1", "type": "lora"}, + "name": "Version 1.0", + "images": [{"url": "https://image.civitai.com/lora1/original=true"}], + "baseModel": "SDXL", + "downloadUrl": "https://civitai.com/lora1/download", + "files": [ + { + "type": "Model", + "primary": True, + "sizeKB": 10240, + "name": "dance_lora_1.safetensors", + "hashes": {"SHA256": "aabbccdd0011"}, + } + ], + } + + lora_info_2 = { + "id": 2398838, + "modelId": 123457, + "model": {"name": "Style LoRA 2", "type": "lora"}, + "name": "Version 2.0", + "images": [{"url": "https://image.civitai.com/lora2/original=true"}], + "baseModel": "SDXL", + "downloadUrl": "https://civitai.com/lora2/download", + "files": [ + { + "type": "Model", + "primary": True, + "sizeKB": 20480, + "name": "style_lora_2.safetensors", + "hashes": {"SHA256": "aabbccdd0022"}, + } + ], + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + if version_id == "2398829": + return lora_info_1, None + if version_id == "2398838": + return lora_info_2, None + return None, "Model not found" + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.civitai_image.get_default_metadata_provider", + fake_metadata_provider, + ) + + parser = CivitaiApiMetadataParser() + + # This simulates the metadata from Civitai image API where modelVersionIds + # is at the root level and meta only contains basic prompt info + metadata = { + "id": 109882763, + "meta": { + "id": 109882763, + "meta": {"prompt": "A woman does the hip bump dance."}, + }, + "modelVersionIds": [2398829, 2398838], + } + + assert parser.is_metadata_matching(metadata) + + result = await parser.parse_metadata(metadata) + + # Verify both LoRAs were created from modelVersionIds + assert len(result["loras"]) == 2 + + # Check first LoRA + lora1 = result["loras"][0] + assert lora1["id"] == 2398829 + assert lora1["name"] == "Dance LoRA 1" + assert lora1["type"] == "lora" + assert lora1["hash"] == "aabbccdd0011" + assert lora1["baseModel"] == "SDXL" + assert ( + lora1["thumbnailUrl"] + == "https://image.civitai.com/lora1/width=450,optimized=true" + ) + + # Check second LoRA + lora2 = result["loras"][1] + assert lora2["id"] == 2398838 + assert lora2["name"] == "Style LoRA 2" + assert lora2["type"] == "lora" + assert lora2["hash"] == "aabbccdd0022" + assert lora2["baseModel"] == "SDXL"