mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-06-10 04:49:24 -03:00
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
This commit is contained in:
@@ -190,25 +190,40 @@ class RecipeEnricher:
|
|||||||
existing_cp = recipe.get("checkpoint")
|
existing_cp = recipe.get("checkpoint")
|
||||||
if existing_cp is None:
|
if existing_cp is None:
|
||||||
existing_cp = {}
|
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)
|
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")
|
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:
|
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"]
|
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:
|
if is_generic and resolved_base_model != current_base_model:
|
||||||
recipe["base_model"] = resolved_base_model
|
recipe["base_model"] = resolved_base_model
|
||||||
|
|
||||||
# 2. Format according to requirements: type, modelId, modelVersionId, modelName, modelVersionName
|
# 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 = {
|
formatted_checkpoint = {
|
||||||
"type": "checkpoint",
|
"type": "checkpoint",
|
||||||
"modelId": checkpoint_data.get("modelId"),
|
"modelId": checkpoint_data.get("modelId"),
|
||||||
"modelVersionId": checkpoint_data.get("id") or checkpoint_data.get("modelVersionId"),
|
"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']
|
"modelName": checkpoint_data.get("name"),
|
||||||
"modelVersionName": checkpoint_data.get("version") # In base.py, 'version' is populated from civitai_data['name']
|
"modelVersionName": checkpoint_data.get("version"),
|
||||||
}
|
}
|
||||||
# Remove None values
|
|
||||||
recipe["checkpoint"] = {k: v for k, v in formatted_checkpoint.items() if v is not None}
|
recipe["checkpoint"] = {k: v for k, v in formatted_checkpoint.items() if v is not None}
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|||||||
@@ -975,6 +975,9 @@ class RecipeManagementHandler:
|
|||||||
civitai_model = civitai_parsed.get("model")
|
civitai_model = civitai_parsed.get("model")
|
||||||
if civitai_model and not metadata.get("checkpoint"):
|
if civitai_model and not metadata.get("checkpoint"):
|
||||||
metadata["checkpoint"] = civitai_model
|
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:
|
elif parsed_embedded:
|
||||||
parsed_loras = parsed_embedded.get("loras")
|
parsed_loras = parsed_embedded.get("loras")
|
||||||
if parsed_loras and not metadata.get("loras"):
|
if parsed_loras and not metadata.get("loras"):
|
||||||
@@ -982,6 +985,8 @@ class RecipeManagementHandler:
|
|||||||
parsed_model = parsed_embedded.get("model")
|
parsed_model = parsed_embedded.get("model")
|
||||||
if parsed_model and not metadata.get("checkpoint"):
|
if parsed_model and not metadata.get("checkpoint"):
|
||||||
metadata["checkpoint"] = parsed_model
|
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()
|
civitai_client = self._civitai_client_getter()
|
||||||
await RecipeEnricher.enrich_recipe(
|
await RecipeEnricher.enrich_recipe(
|
||||||
@@ -1489,13 +1494,16 @@ class RecipeManagementHandler:
|
|||||||
if not image_url:
|
if not image_url:
|
||||||
raise RecipeValidationError("Missing required field: 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)
|
image_id = extract_civitai_image_id(image_url)
|
||||||
if not image_id:
|
if not image_id:
|
||||||
raise RecipeValidationError(
|
raise RecipeValidationError(
|
||||||
"Could not extract Civitai image ID from URL"
|
"Could not extract Civitai image ID from URL"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check for duplicate (fast, before acquiring semaphore)
|
# Check for duplicate (fast, before acquiring semaphore), unless force
|
||||||
|
if not force:
|
||||||
cache = await recipe_scanner.get_cached_data()
|
cache = await recipe_scanner.get_cached_data()
|
||||||
for recipe in getattr(cache, "raw_data", []):
|
for recipe in getattr(cache, "raw_data", []):
|
||||||
source = recipe.get("source_path")
|
source = recipe.get("source_path")
|
||||||
@@ -1613,6 +1621,9 @@ class RecipeManagementHandler:
|
|||||||
civitai_model = civitai_parsed.get("model")
|
civitai_model = civitai_parsed.get("model")
|
||||||
if civitai_model and not metadata.get("checkpoint"):
|
if civitai_model and not metadata.get("checkpoint"):
|
||||||
metadata["checkpoint"] = civitai_model
|
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:
|
elif parsed_embedded:
|
||||||
parsed_loras = parsed_embedded.get("loras")
|
parsed_loras = parsed_embedded.get("loras")
|
||||||
if parsed_loras and not metadata.get("loras"):
|
if parsed_loras and not metadata.get("loras"):
|
||||||
@@ -1620,6 +1631,8 @@ class RecipeManagementHandler:
|
|||||||
parsed_model = parsed_embedded.get("model")
|
parsed_model = parsed_embedded.get("model")
|
||||||
if parsed_model and not metadata.get("checkpoint"):
|
if parsed_model and not metadata.get("checkpoint"):
|
||||||
metadata["checkpoint"] = parsed_model
|
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()
|
civitai_client = self._civitai_client_getter()
|
||||||
await RecipeEnricher.enrich_recipe(
|
await RecipeEnricher.enrich_recipe(
|
||||||
|
|||||||
Reference in New Issue
Block a user