From a552f07448e7be2b38f57423179bb45adb377a37 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Wed, 24 Dec 2025 20:25:39 +0800 Subject: [PATCH] feat: Refactor checkpoint metadata to use Civitai API naming conventions and remove `gen_params` checkpoint syncing. --- py/recipes/enrichment.py | 36 ++++----- py/recipes/merger.py | 2 +- py/services/recipe_scanner.py | 2 +- refs/recipe.json | 99 ++++++------------------ tests/services/test_gen_params_merger.py | 3 +- tests/services/test_recipe_repair.py | 28 ++++--- 6 files changed, 62 insertions(+), 108 deletions(-) diff --git a/py/recipes/enrichment.py b/py/recipes/enrichment.py index 28b24430..83274164 100644 --- 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 diff --git a/py/recipes/merger.py b/py/recipes/merger.py index b0702656..93d19857 100644 --- 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 = { diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index 6cd80a0e..8b7554ff 100644 --- 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, diff --git a/refs/recipe.json b/refs/recipe.json index bcf1d6b8..12630d33 100644 --- 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" } \ No newline at end of file diff --git a/tests/services/test_gen_params_merger.py b/tests/services/test_gen_params_merger.py index cacfe28d..16313f0a 100644 --- 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 = { diff --git a/tests/services/test_recipe_repair.py b/tests/services/test_recipe_repair.py index 6b9bc6eb..bc1020cd 100644 --- 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"