fix(recipe): use resources type field to identify checkpoint instead of modelVersionIds[0]

When importing a CivitAI image as a recipe, modelVersionIds[0] was blindly used as the checkpoint version ID. This array mixes checkpoints and LoRAs without ordering guarantees, causing LoRAs to be saved as the recipe checkpoint.

Fix by:
1. Removing the modelVersionIds[0] fallback in _download_remote_media
2. Parsing resources entries with type:"model" as the checkpoint
3. Adding model type validation in populate_checkpoint_from_civitai

Also add 2 tests for the new behavior and fix 3 tests whose mocks lacked the required model.type field.
This commit is contained in:
Will Miao
2026-05-28 15:46:33 +08:00
parent 3f6824eef6
commit 34791c2ad7
6 changed files with 203 additions and 10 deletions

View File

@@ -467,7 +467,10 @@ async def test_import_remote_recipe(monkeypatch, tmp_path: Path) -> None:
class Provider:
async def get_model_version_info(self, model_version_id):
provider_calls.append(model_version_id)
return {"baseModel": "Flux Provider"}, None
return {
"baseModel": "Flux Provider",
"model": {"type": "Checkpoint", "name": "Flux"},
}, None
async def fake_get_default_metadata_provider():
return Provider()

View File

@@ -298,3 +298,113 @@ async def test_parse_metadata_handles_modelVersionIds(monkeypatch):
assert lora2["type"] == "lora"
assert lora2["hash"] == "aabbccdd0022"
assert lora2["baseModel"] == "SDXL"
@pytest.mark.asyncio
async def test_parse_metadata_extracts_checkpoint_from_resources_model_type(monkeypatch):
"""resources entries with type:"model" should be captured as the checkpoint,
not skipped (which was the old buggy behavior), and not mixed into loras."""
captured_hashes = []
async def fake_metadata_provider():
class Provider:
async def get_model_by_hash(self, model_hash):
captured_hashes.append(model_hash)
if model_hash == "a1b2c3d4e5":
return ({
"id": 999,
"modelId": 888,
"name": "v1.0",
"model": {"name": "Real Checkpoint", "type": "Checkpoint"},
"baseModel": "SDXL 1.0",
"images": [{"url": "https://image.civitai.com/cp/original=true"}],
"files": [{"type": "Model", "primary": True, "sizeKB": 1024, "name": "cp.safetensors"}]
}, 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",
"resources": [
{"hash": "a1b2c3d4e5", "name": "Real Checkpoint", "type": "model"},
{"hash": "f6g7h8i9j0", "name": "Some LoRA", "type": "lora", "weight": 0.8},
],
"Model hash": "a1b2c3d4e5",
}
result = await parser.parse_metadata(metadata)
# The type:"model" resource should be in result["model"], not in result["loras"]
assert result["model"] is not None, "checkpoint model should be extracted"
assert result["model"]["name"] == "Real Checkpoint"
assert result["model"]["hash"] == "a1b2c3d4e5"
assert result["model"]["type"] == "model"
# The LoRA resource should be in result["loras"]
assert len(result["loras"]) == 1
assert result["loras"][0]["name"] == "Some LoRA"
# The checkpoint hash should have triggered a lookup
assert "a1b2c3d4e5" in captured_hashes
@pytest.mark.asyncio
async def test_parse_metadata_resources_model_type_does_not_duplicate_checkpoint_in_loras(monkeypatch):
"""When a resources entry has type:"model", it should NOT also appear in loras.
Regression test for the bug where the checkpoint model appeared in both places."""
async def fake_metadata_provider():
class Provider:
async def get_model_by_hash(self, model_hash):
if model_hash == "cp123hash":
return ({
"id": 100,
"modelId": 200,
"name": "v2",
"model": {"name": "My Checkpoint", "type": "Checkpoint"},
"baseModel": "SDXL",
"files": [{"type": "Model", "primary": True, "sizeKB": 1024, "name": "cp.safetensors"}]
}, None)
if model_hash == "lora1hash":
return ({
"id": 300,
"modelId": 400,
"name": "v1",
"model": {"name": "Style LoRA", "type": "LORA"},
"baseModel": "SDXL",
"files": [{"type": "Model", "primary": True, "sizeKB": 512, "name": "style.safetensors"}]
}, 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 = {
"resources": [
{"hash": "cp123hash", "name": "My Checkpoint", "type": "model"},
{"hash": "lora1hash", "name": "Style LoRA", "type": "lora", "weight": 0.5},
],
}
result = await parser.parse_metadata(metadata)
# Checkpoint must NOT appear in loras
lora_names = {l["name"] for l in result["loras"]}
assert "My Checkpoint" not in lora_names
assert "Style LoRA" in lora_names
# Checkpoint must be in result["model"]
assert result["model"] is not None
assert result["model"]["name"] == "My Checkpoint"

View File

@@ -94,7 +94,7 @@ async def test_repair_all_recipes_with_enriched_checkpoint_id(setup_scanner):
"id": 5678,
"modelId": 1234,
"name": "v1.0",
"model": {"name": "Full Model Name"},
"model": {"name": "Full Model Name", "type": "Checkpoint"},
"baseModel": "SDXL 1.0",
"images": [{"url": "https://image.url/thumb.jpg"}],
"files": [{"type": "Model", "hashes": {"SHA256": "ABCDEF"}, "name": "full_filename.safetensors"}]
@@ -142,7 +142,7 @@ async def test_repair_all_recipes_supports_civitai_red_source_url(setup_scanner)
"id": 5678,
"modelId": 1234,
"name": "v1.0",
"model": {"name": "Full Model Name"},
"model": {"name": "Full Model Name", "type": "Checkpoint"},
"baseModel": "SDXL 1.0",
"images": [{"url": "https://image.url/thumb.jpg"}],
"files": [
@@ -183,7 +183,7 @@ async def test_repair_all_recipes_with_enriched_checkpoint_hash(setup_scanner):
"id": 999,
"modelId": 888,
"name": "v2.0",
"model": {"name": "Hashed Model"},
"model": {"name": "Hashed Model", "type": "Checkpoint"},
"baseModel": "SD 1.5",
"files": [{"type": "Model", "hashes": {"SHA256": "hash123"}, "name": "hashed.safetensors"}]
}, None)