feat: Refactor checkpoint metadata to use Civitai API naming conventions and remove gen_params checkpoint syncing.

2026-03-21 21:22:11 -03:00 · 2025-12-24 20:25:39 +08:00
parent 6486107ca2
commit a552f07448
6 changed files with 62 additions and 108 deletions
--- a/py/recipes/enrichment.py
+++ b/py/recipes/enrichment.py
@@ -124,21 +124,10 @@ class RecipeEnricher:
                recipe, target_version_id, target_hash, model_val, checkpoint_val
            )
            if checkpoint_updated:
-                # Sync to gen_params for consistency with legacy usage
-                if "gen_params" not in recipe:
-                    recipe["gen_params"] = {}
-                recipe["gen_params"]["checkpoint"] = recipe["checkpoint"]
                updated = True
        else:
-            # Even if we have a checkpoint, ensure it is synced to gen_params if missing there
-            if "checkpoint" in recipe and recipe["checkpoint"]:
-                if "gen_params" not in recipe:
-                    recipe["gen_params"] = {}
-                if "checkpoint" not in recipe["gen_params"]:
-                    recipe["gen_params"]["checkpoint"] = recipe["checkpoint"]
-                    # We don't necessarily mark 'updated=True' just for this sync if the rest is the same,
-                    # but it's safer to ensure it's there.
-                    updated = True
+            # Checkpoint exists, no need to sync to gen_params anymore.
+            pass
        # If base_model is empty or very generic, try to use what we found in checkpoint
        current_base_model = recipe.get("base_model")
        checkpoint_after = recipe.get("checkpoint")
@@ -201,23 +190,28 @@ class RecipeEnricher:
            if existing_cp is None:
                existing_cp = {}
            checkpoint_data = await RecipeMetadataParser.populate_checkpoint_from_civitai(existing_cp, civitai_info)
-            recipe["checkpoint"] = checkpoint_data
            
-            # Ensure the modelVersionId is stored if we found it
-            if target_version_id and "modelVersionId" not in recipe["checkpoint"]:
-                recipe["checkpoint"]["modelVersionId"] = int(target_version_id)
+            # Format according to requirements: type, modelId, modelVersionId, modelName, modelVersionName
+            formatted_checkpoint = {
+                "type": "checkpoint",
+                "modelId": checkpoint_data.get("modelId"),
+                "modelVersionId": checkpoint_data.get("id") or checkpoint_data.get("modelVersionId"),
+                "modelName": checkpoint_data.get("name"), # In base.py, 'name' is populated from civitai_data['model']['name']
+                "modelVersionName": checkpoint_data.get("version") # In base.py, 'version' is populated from civitai_data['name']
+            }
+            # Remove None values
+            recipe["checkpoint"] = {k: v for k, v in formatted_checkpoint.items() if v is not None}
+            
            return True
        else:
            # Fallback to name extraction if we don't already have one
            existing_cp = recipe.get("checkpoint")
-            if not existing_cp or not existing_cp.get("name"):
+            if not existing_cp or not existing_cp.get("modelName"):
                cp_name = checkpoint_val
                if cp_name:
                    recipe["checkpoint"] = {
                        "type": "checkpoint",
-                        "name": cp_name,
-                        "modelName": cp_name,
-                        "file_name": os.path.splitext(cp_name)[0]
+                        "modelName": cp_name
                    }
                    return True
                
--- a/py/recipes/merger.py
+++ b/py/recipes/merger.py
@@ -12,7 +12,7 @@ class GenParamsMerger:
        "baseModel", "resources", "disablePoi", "aspectRatio", "Created Date",
        "experimental", "civitaiResources", "civitai_resources", "Civitai resources",
        "modelVersionId", "modelId", "hashes", "Model", "Model hash", "checkpoint_hash",
-        "checksum", "model_checksum"
+        "checkpoint", "checksum", "model_checksum"
    }
    
    NORMALIZATION_MAPPING = {
--- a/py/services/recipe_scanner.py
+++ b/py/services/recipe_scanner.py
@@ -57,7 +57,7 @@ class RecipeScanner:
            cls._instance._civitai_client = None  # Will be lazily initialized
        return cls._instance
    
-    REPAIR_VERSION = 2
+    REPAIR_VERSION = 3
    
    def __init__(
        self,
--- a/refs/recipe.json
+++ b/refs/recipe.json
@@ -1,82 +1,33 @@
 {
-    "id": "0448c06d-de1b-46ab-975c-c5aa60d90dbc",
-    "file_path": "D:/Workspace/ComfyUI/models/loras/recipes/0448c06d-de1b-46ab-975c-c5aa60d90dbc.jpg",
-    "title": "a mysterious, steampunk-inspired character standing in a dramatic pose",
-    "modified": 1741837612.3931093,
-    "created_date": 1741492786.5581934,
-    "base_model": "Flux.1 D",
+    "id": "42803a29-02dc-49e1-b798-27da70e8b408",
+    "file_path": "/home/miao/workspace/ComfyUI/models/loras/recipes/test/42803a29-02dc-49e1-b798-27da70e8b408.webp",
+    "title": "masterpiece, best quality, amazing quality, very aesthetic, detailed eyes, perfect",
+    "modified": 1754897325.0507245,
+    "created_date": 1754897325.0507245,
+    "base_model": "Illustrious",
    "loras": [
        {
-            "file_name": "ChronoDivinitiesFlux_r1",
-            "hash": "ddbc5abd00db46ad464f5e3ca85f8f7121bc14b594d6785f441d9b002fffe66a",
-            "strength": 0.8,
-            "modelVersionId": 1438879,
-            "modelName": "Chrono Divinities - By HailoKnight",
-            "modelVersionName": "Flux"
-        },
-        {
-            "file_name": "flux.1_lora_flyway_ink-dynamic",
-            "hash": "4b4f3b469a0d5d3a04a46886abfa33daa37a905db070ccfbd10b345c6fb00eff",
-            "strength": 0.2,
-            "modelVersionId": 914935,
-            "modelName": "Ink-style",
-            "modelVersionName": "ink-dynamic"
-        },
-        {
-            "file_name": "ck-painterly-fantasy-000017",
-            "hash": "48c67064e2936aec342580a2a729d91d75eb818e45ecf993b9650cc66c94c420",
-            "strength": 0.2,
-            "modelVersionId": 1189379,
-            "modelName": "Painterly Fantasy by ChronoKnight - [FLUX & IL]",
-            "modelVersionName": "FLUX"
-        },
-        {
-            "file_name": "RetroAnimeFluxV1",
-            "hash": "8f43c31b6c3238ac44195c970d511d759c5893bddd00f59f42b8fe51e8e76fa0",
-            "strength": 0.8,
-            "modelVersionId": 806265,
-            "modelName": "Retro Anime Flux - Style",
-            "modelVersionName": "v1.0"
-        },
-        {
-            "file_name": "Mezzotint_Artstyle_for_Flux_-_by_Ethanar",
-            "hash": "e6961502769123bf23a66c5c5298d76264fd6b9610f018319a0ccb091bfc308e",
-            "strength": 0.2,
-            "modelVersionId": 757030,
-            "modelName": "Mezzotint Artstyle for Flux - by Ethanar",
-            "modelVersionName": "V1"
-        },
-        {
-            "file_name": "FluxMythG0thicL1nes",
-            "hash": "ecb03595de62bd6183a0dd2b38bea35669fd4d509f4bbae5aa0572cfb7ef4279",
-            "strength": 0.4,
-            "modelVersionId": 1202162,
-            "modelName": "Velvet's Mythic Fantasy Styles | Flux + Pony + illustrious",
-            "modelVersionName": "Flux Gothic Lines"
-        },
-        {
-            "file_name": "Elden_Ring_-_Yoshitaka_Amano",
-            "hash": "c660c4c55320be7206cb6a917c59d8da3953cc07169fe10bda833a54ec0024f9",
-            "strength": 0.75,
-            "modelVersionId": 746484,
-            "modelName": "Elden Ring - Yoshitaka Amano",
-            "modelVersionName": "V1"
+            "file_name": "",
+            "hash": "1b5b763d83961bb5745f3af8271ba83f1d4fd69c16278dae6d5b4e194bdde97a",
+            "strength": 1.0,
+            "modelVersionId": 2007092,
+            "modelName": "Pony: People's Works +",
+            "modelVersionName": "v8_Illusv1.0",
+            "isDeleted": false,
+            "exclude": false
        }
    ],
    "gen_params": {
-        "prompt": "a mysterious, steampunk-inspired character standing in a dramatic pose. The character is dressed in a long, intricately detailed dark coat with ornate patterns, a wide-brimmed hat, and leather boots. The face is partially obscured by the hat's shadow, adding to the enigmatic aura. The background showcases a large, antique clock with Roman numerals, surrounded by dynamic lightning and ethereal white birds, enhancing the fantastical atmosphere. The color palette is dominated by dark tones with striking contrasts of white and blue lightning, creating a sense of tension and energy. The overall composition is vertical, with the character centrally positioned, exuding a sense of power and mystery. hkchrono",
-        "negative_prompt": "",
-        "checkpoint": {
-            "type": "checkpoint",
-            "modelVersionId": 691639,
-            "modelName": "FLUX",
-            "modelVersionName": "Dev"
-        },
-        "steps": "30",
-        "sampler": "Undefined",
-        "cfg_scale": "3.5",
-        "seed": "1472903449",
+        "prompt": "masterpiece, best quality, amazing quality, very aesthetic, detailed eyes, perfect eyes, realistic eyes,\n(flat colors:1.5), (anime:1.5), (lineart:1.5),\nclose-up, solo, tongue, 1girl, food, (saliva:0.1), open mouth, candy, simple background, blue background, large lollipop, tongue out, fade background, lips, hand up, holding, looking at viewer, licking, seductive, half-closed eyes,",
+        "negative_prompt": "shiny skin,",
+        "steps": 19,
+        "sampler": "Euler a",
+        "cfg_scale": 5,
+        "seed": 1765271748,
        "size": "832x1216",
-        "clip_skip": "2"
-    }
+        "clip_skip": 2
+    },
+    "fingerprint": "1b5b763d83961bb5745f3af8271ba83f1d4fd69c16278dae6d5b4e194bdde97a:1.0",
+    "source_path": "https://civitai.com/images/92427432",
+    "folder": "test"
 }
--- a/tests/services/test_gen_params_merger.py
+++ b/tests/services/test_gen_params_merger.py
@@ -45,7 +45,7 @@ def test_merge_none_values():
    assert merged == {}

 def test_merge_filters_blacklisted_keys():
-    request_params = {"prompt": "test", "id": "should-be-removed"}
+    request_params = {"prompt": "test", "id": "should-be-removed", "checkpoint": "should-not-be-here"}
    civitai_meta = {"cfg": 7, "url": "remove-me"}
    embedded_metadata = {"seed": 123, "hash": "remove-also"}
    
@@ -57,6 +57,7 @@ def test_merge_filters_blacklisted_keys():
    assert "id" not in merged
    assert "url" not in merged
    assert "hash" not in merged
+    assert "checkpoint" not in merged

 def test_merge_filters_meta_and_normalizes_keys():
    civitai_meta = {
--- a/tests/services/test_recipe_repair.py
+++ b/tests/services/test_recipe_repair.py
@@ -109,13 +109,15 @@ async def test_repair_all_recipes_with_enriched_checkpoint_id(setup_scanner):
    
    saved_recipe = recipe_scanner._save_recipe_persistently.call_args[0][0]
    checkpoint = saved_recipe["checkpoint"]
-    assert checkpoint["name"] == "Full Model Name"
-    assert checkpoint["version"] == "v1.0"
+    assert checkpoint["modelName"] == "Full Model Name"
+    assert checkpoint["modelVersionName"] == "v1.0"
    assert checkpoint["modelId"] == 1234
-    assert checkpoint["id"] == 5678
-    assert checkpoint["hash"] == "abcdef"
-    assert checkpoint["file_name"] == "full_filename"
-    assert "thumbnailUrl" not in checkpoint # Stripped during sanitation
+    assert checkpoint["modelVersionId"] == 5678
+    assert checkpoint["type"] == "checkpoint"
+    assert "name" not in checkpoint
+    assert "version" not in checkpoint
+    assert "hash" not in checkpoint
+    assert "file_name" not in checkpoint

@pytest.mark.asyncio
 async def test_repair_all_recipes_with_enriched_checkpoint_hash(setup_scanner):
@@ -151,10 +153,10 @@ async def test_repair_all_recipes_with_enriched_checkpoint_hash(setup_scanner):
    
    saved_recipe = recipe_scanner._save_recipe_persistently.call_args[0][0]
    checkpoint = saved_recipe["checkpoint"]
-    assert checkpoint["name"] == "Hashed Model"
-    assert checkpoint["version"] == "v2.0"
+    assert checkpoint["modelName"] == "Hashed Model"
+    assert checkpoint["modelVersionName"] == "v2.0"
    assert checkpoint["modelId"] == 888
-    assert checkpoint["hash"] == "hash123"
+    assert checkpoint["type"] == "checkpoint"

@pytest.mark.asyncio
 async def test_repair_all_recipes_fallback_to_basic(setup_scanner):
@@ -180,7 +182,8 @@ async def test_repair_all_recipes_fallback_to_basic(setup_scanner):
    # Verify
    assert results["repaired"] == 1
    saved_recipe = recipe_scanner._save_recipe_persistently.call_args[0][0]
-    assert saved_recipe["checkpoint"]["name"] == "just_a_name.safetensors"
+    assert saved_recipe["checkpoint"]["modelName"] == "just_a_name.safetensors"
+    assert saved_recipe["checkpoint"]["type"] == "checkpoint"
    assert "modelId" not in saved_recipe["checkpoint"]

@pytest.mark.asyncio
@@ -271,4 +274,9 @@ async def test_sanitize_recipe_for_storage(recipe_scanner):
    assert "strength" in clean["loras"][0]
    assert clean["loras"][0]["strength"] == 0.5
    assert "localPath" not in clean["checkpoint"]
+    # Testing based on what enricher would produce if it ran, 
+    # but here we are just testing the sanitizer which handles what is ALREADY there.
+    # However, the sanitizer doesn't rename fields, it just removes runtime ones.
+    # Since we changed the enricher to NOT put 'name' anymore, this test case 
+    # should probably reflect the new fields if it's simulating a real recipe.
    assert clean["checkpoint"]["name"] == "CP"