From ccf1c6f2ae43fea5e11f932091fc53c2e635bb70 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 1 Jun 2026 17:58:08 +0800 Subject: [PATCH] fix(recipe): resolve base_model from parser and prevent empty checkpoint save on CivitAI import - Apply CivitaiApiMetadataParser's base_model result to metadata in _do_import_remote_recipe and _do_import_from_url (was previously discarded) - Extract baseModel from raw civitai_info before populate_checkpoint_from_civitai so it's not lost when the type check rejects non-checkpoint model versions - Only format and save checkpoint entry when it has real data (modelId, versionId, name, or version), preventing empty {'type': 'checkpoint'} stubs --- py/recipes/enrichment.py | 45 ++++++++++++++++++--------- py/routes/handlers/recipe_handlers.py | 39 +++++++++++++++-------- 2 files changed, 56 insertions(+), 28 deletions(-) diff --git a/py/recipes/enrichment.py b/py/recipes/enrichment.py index f640bb32..db1475af 100644 --- a/py/recipes/enrichment.py +++ b/py/recipes/enrichment.py @@ -190,27 +190,42 @@ class RecipeEnricher: existing_cp = recipe.get("checkpoint") if existing_cp is None: existing_cp = {} + + # Extract baseModel from raw civitai_info before populate_checkpoint_from_civitai + # (populate may reject non-checkpoint types and lose this data) + base_model_from_civitai: str = "" + if isinstance(civitai_info, dict): + base_model_from_civitai = civitai_info.get("baseModel", "") or "" + elif isinstance(civitai_info, tuple) and len(civitai_info) > 0 and isinstance(civitai_info[0], dict): + base_model_from_civitai = civitai_info[0].get("baseModel", "") or "" + checkpoint_data = await RecipeMetadataParser.populate_checkpoint_from_civitai(existing_cp, civitai_info) - # 1. First, resolve base_model using full data before we format it away + + # 1. Resolve base_model from checkpoint_data first, then fall back to raw civitai_info current_base_model = recipe.get("base_model") - resolved_base_model = checkpoint_data.get("baseModel") + resolved_base_model = checkpoint_data.get("baseModel") or base_model_from_civitai if resolved_base_model: - # Update if empty OR if it matches our generic prefix but is less specific is_generic = not current_base_model or current_base_model.lower() in ["flux", "sdxl", "sd15"] if is_generic and resolved_base_model != current_base_model: recipe["base_model"] = resolved_base_model - - # 2. 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} - + + # 2. Only format and save checkpoint if it has real data (not just type after type rejection) + has_checkpoint_data = any([ + checkpoint_data.get("modelId"), + checkpoint_data.get("id") or checkpoint_data.get("modelVersionId"), + checkpoint_data.get("name"), + checkpoint_data.get("version"), + ]) + if has_checkpoint_data: + formatted_checkpoint = { + "type": "checkpoint", + "modelId": checkpoint_data.get("modelId"), + "modelVersionId": checkpoint_data.get("id") or checkpoint_data.get("modelVersionId"), + "modelName": checkpoint_data.get("name"), + "modelVersionName": checkpoint_data.get("version"), + } + 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 diff --git a/py/routes/handlers/recipe_handlers.py b/py/routes/handlers/recipe_handlers.py index 413b6eab..7d0574a0 100644 --- a/py/routes/handlers/recipe_handlers.py +++ b/py/routes/handlers/recipe_handlers.py @@ -975,6 +975,9 @@ class RecipeManagementHandler: civitai_model = civitai_parsed.get("model") if civitai_model and not metadata.get("checkpoint"): metadata["checkpoint"] = civitai_model + civitai_base_model = civitai_parsed.get("base_model") + if civitai_base_model and not metadata.get("base_model"): + metadata["base_model"] = civitai_base_model elif parsed_embedded: parsed_loras = parsed_embedded.get("loras") if parsed_loras and not metadata.get("loras"): @@ -982,6 +985,8 @@ class RecipeManagementHandler: parsed_model = parsed_embedded.get("model") if parsed_model and not metadata.get("checkpoint"): metadata["checkpoint"] = parsed_model + if parsed_embedded.get("base_model") and not metadata.get("base_model"): + metadata["base_model"] = parsed_embedded["base_model"] civitai_client = self._civitai_client_getter() await RecipeEnricher.enrich_recipe( @@ -1489,25 +1494,28 @@ class RecipeManagementHandler: if not image_url: raise RecipeValidationError("Missing required field: image_url") + force = request.query.get("force", "false").lower() == "true" + image_id = extract_civitai_image_id(image_url) if not image_id: raise RecipeValidationError( "Could not extract Civitai image ID from URL" ) - # Check for duplicate (fast, before acquiring semaphore) - cache = await recipe_scanner.get_cached_data() - for recipe in getattr(cache, "raw_data", []): - source = recipe.get("source_path") - if source: - existing_id = extract_civitai_image_id(source) - if existing_id == image_id: - return web.json_response({ - "success": True, - "recipe_id": recipe.get("id"), - "name": recipe.get("title", ""), - "already_exists": True, - }) + # Check for duplicate (fast, before acquiring semaphore), unless force + if not force: + cache = await recipe_scanner.get_cached_data() + for recipe in getattr(cache, "raw_data", []): + source = recipe.get("source_path") + if source: + existing_id = extract_civitai_image_id(source) + if existing_id == image_id: + return web.json_response({ + "success": True, + "recipe_id": recipe.get("id"), + "name": recipe.get("title", ""), + "already_exists": True, + }) async with self._import_semaphore: return await self._do_import_from_url(image_url, recipe_scanner) @@ -1613,6 +1621,9 @@ class RecipeManagementHandler: civitai_model = civitai_parsed.get("model") if civitai_model and not metadata.get("checkpoint"): metadata["checkpoint"] = civitai_model + civitai_base_model = civitai_parsed.get("base_model") + if civitai_base_model and not metadata.get("base_model"): + metadata["base_model"] = civitai_base_model elif parsed_embedded: parsed_loras = parsed_embedded.get("loras") if parsed_loras and not metadata.get("loras"): @@ -1620,6 +1631,8 @@ class RecipeManagementHandler: parsed_model = parsed_embedded.get("model") if parsed_model and not metadata.get("checkpoint"): metadata["checkpoint"] = parsed_model + if parsed_embedded.get("base_model") and not metadata.get("base_model"): + metadata["base_model"] = parsed_embedded["base_model"] civitai_client = self._civitai_client_getter() await RecipeEnricher.enrich_recipe(