import pytest from py.recipes.parsers.civitai_image import CivitaiApiMetadataParser @pytest.mark.asyncio async def test_parse_metadata_creates_loras_from_hashes(monkeypatch): async def fake_metadata_provider(): return None monkeypatch.setattr( "py.recipes.parsers.civitai_image.get_default_metadata_provider", fake_metadata_provider, ) parser = CivitaiApiMetadataParser() metadata = { "Size": "1536x2688", "seed": 3766932689, "Model": "indexed_v1", "steps": 30, "hashes": { "model": "692186a14a", "LORA:Jedst1": "fb4063c470", "LORA:HassaKu_style": "3ce00b926b", "LORA:DetailedEyes_V3": "2c1c3f889f", "LORA:jiaocha_illustriousXL": "35d3e6f8b0", "LORA:绪儿 厚涂构图光影质感增强V3": "d9b5900a59", }, "prompt": "test", "Version": "ComfyUI", "sampler": "er_sde_ays_30", "cfgScale": 5, "clipSkip": 2, "resources": [ { "hash": "692186a14a", "name": "indexed_v1", "type": "model", } ], "Model hash": "692186a14a", "negativePrompt": "bad", "username": "LumaRift", "baseModel": "Illustrious", } result = await parser.parse_metadata(metadata) assert result["base_model"] == "Illustrious" assert len(result["loras"]) == 5 assert all(lora["weight"] == 1.0 for lora in result["loras"]) assert {lora["name"] for lora in result["loras"]} == { "Jedst1", "HassaKu_style", "DetailedEyes_V3", "jiaocha_illustriousXL", "绪儿 厚涂构图光影质感增强V3", } @pytest.mark.asyncio async def test_parse_metadata_handles_nested_meta_and_lowercase_hashes(monkeypatch): async def fake_metadata_provider(): return None monkeypatch.setattr( "py.recipes.parsers.civitai_image.get_default_metadata_provider", fake_metadata_provider, ) parser = CivitaiApiMetadataParser() metadata = { "id": 106706587, "meta": { "prompt": "An enigmatic silhouette", "hashes": { "model": "ee75fd24a4", "lora:mj": "de49e1e98c", "LORA:Another_Earth_2": "dc11b64a8b", }, "resources": [ { "hash": "ee75fd24a4", "name": "stoiqoNewrealityFLUXSD35_f1DAlphaTwo", "type": "model", } ], }, } assert parser.is_metadata_matching(metadata) result = await parser.parse_metadata(metadata) assert result["gen_params"]["prompt"] == "An enigmatic silhouette" assert {l["name"] for l in result["loras"]} == {"mj", "Another_Earth_2"} assert {l["hash"] for l in result["loras"]} == {"de49e1e98c", "dc11b64a8b"} @pytest.mark.asyncio async def test_parse_metadata_populates_checkpoint_and_rewrites_thumbnails(monkeypatch): checkpoint_info = { "id": 222, "modelId": 111, "model": {"name": "Checkpoint Example", "type": "checkpoint"}, "name": "Checkpoint Version", "images": [{"url": "https://image.civitai.com/checkpoints/original=true"}], "baseModel": "Illustrious", "downloadUrl": "https://civitai.com/checkpoint/download", "files": [ { "type": "Model", "primary": True, "sizeKB": 1024, "name": "Checkpoint Example.safetensors", "hashes": {"SHA256": "FFAA0011"}, } ], } lora_info = { "id": 444, "modelId": 333, "model": {"name": "Example Lora Model", "type": "lora"}, "name": "Example Lora Version", "images": [{"url": "https://image.civitai.com/loras/original=true"}], "baseModel": "Illustrious", "downloadUrl": "https://civitai.com/lora/download", "files": [ { "type": "Model", "primary": True, "sizeKB": 512, "hashes": {"SHA256": "abc123"}, } ], } async def fake_metadata_provider(): class Provider: async def get_model_version_info(self, version_id): if version_id == "222": return checkpoint_info, None if version_id == "444": return lora_info, None return None, "Model not found" return Provider() monkeypatch.setattr( "py.recipes.parsers.civitai_image.get_default_metadata_provider", fake_metadata_provider, ) parser = CivitaiApiMetadataParser() metadata = { "prompt": "test prompt", "negativePrompt": "test negative prompt", "civitaiResources": [ { "type": "checkpoint", "modelId": 111, "modelVersionId": 222, "modelName": "Checkpoint Example", "modelVersionName": "Checkpoint Version", }, { "type": "lora", "modelId": 333, "modelVersionId": 444, "modelName": "Example Lora", "modelVersionName": "Lora Version", "weight": 0.7, }, ], } result = await parser.parse_metadata(metadata) 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"]["modelId"] == 111 assert result["model"]["size"] == 1024 * 1024 assert result["model"]["hash"] == "ffaa0011" assert result["model"]["file_name"] == "Checkpoint Example" 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]["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"