alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-06-09 20:39:25 -03:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

16
py/recipes/base.py
View File

@@ -7,7 +7,7 @@ import re
from typing import Dict, List, Any, Optional, Tuple
from abc import ABC, abstractmethod
from ..config import config
from ..utils.constants import VALID_LORA_TYPES
from ..utils.constants import VALID_LORA_TYPES, VALID_CHECKPOINT_SUB_TYPES
from ..utils.civitai_utils import rewrite_preview_url
logger = logging.getLogger(__name__)
@@ -173,6 +173,20 @@ class RecipeMetadataParser(ABC):
checkpoint['isDeleted'] = True
return checkpoint
# Validate that the model type is actually a checkpoint.
# Unlike populate_lora_from_civitai which has this check,
# this function was missing type validation — allowing LoRA
# version data to be saved as the recipe's checkpoint when the
# wrong version ID was passed downstream (fixed in v2.7+).
model_type = civitai_data.get('model', {}).get('type', '').lower()
if model_type not in VALID_CHECKPOINT_SUB_TYPES:
logger.warning(
f"Cannot populate checkpoint: model version {civitai_data.get('id')} "
f"has type '{model_type}', expected one of {VALID_CHECKPOINT_SUB_TYPES}. "
f"Skipping checkpoint enrichment."
)
return checkpoint
if 'model' in civitai_data and 'name' in civitai_data['model']:
checkpoint['name'] = civitai_data['model']['name']

61
py/recipes/parsers/civitai_image.py
View File

@@ -185,8 +185,67 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Process standard resources array
if "resources" in metadata and isinstance(metadata["resources"], list):
for resource in metadata["resources"]:
resource_type = resource.get("type", "lora")
# Track resources with type "model" — these are checkpoint models.
# The resources array is the most reliable source for checkpoint
# identification because it has an explicit type field and hash,
# unlike modelVersionIds which is a flat list with no type info.
if resource_type == "model":
checkpoint_entry = {
"id": 0,
"modelId": 0,
"name": resource.get("name", "Unknown Model"),
"version": "",
"type": resource.get("type", "model"),
"existsLocally": False,
"localPath": None,
"file_name": resource.get("name", ""),
"hash": resource.get("hash", "") or "",
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# Try to look up base model from the checkpoint hash
if checkpoint_entry["hash"] and metadata_provider:
try:
civitai_info = (
await metadata_provider.get_model_by_hash(
checkpoint_entry["hash"]
)
)
civitai_data, error_msg = (
(civitai_info, None)
if not isinstance(civitai_info, tuple)
else civitai_info
)
if civitai_data and error_msg != "Model not found":
if 'model' in civitai_data and 'name' in civitai_data['model']:
checkpoint_entry['name'] = civitai_data['model']['name']
checkpoint_entry['id'] = civitai_data.get('id', 0)
checkpoint_entry['modelId'] = civitai_data.get('modelId', 0)
if 'name' in civitai_data:
checkpoint_entry['version'] = civitai_data['name']
base_model = civitai_data.get('baseModel', '')
if base_model:
checkpoint_entry['baseModel'] = base_model
if not result['base_model']:
result['base_model'] = base_model
except Exception as e:
logger.error(
f"Error fetching checkpoint info for hash "
f"{checkpoint_entry['hash']}: {e}"
)
if result["model"] is None:
result["model"] = checkpoint_entry
continue
# Modified to process resources without a type field as potential LoRAs
if resource.get("type", "lora") == "lora":
if resource_type == "lora":
lora_hash = resource.get("hash", "")
# Try to get hash from the hashes field if not present in resource

15
py/routes/handlers/recipe_handlers.py
View File

@@ -1293,11 +1293,18 @@ class RecipeManagementHandler:
image_info.get("meta") if civitai_image_id and image_info else None
)
if civitai_image_id and image_info:
# modelVersionId (singular) — the primary version for this
# image on CivitAI. May be absent, or may *not* be the
# checkpoint (e.g. when the image was generated with a LoRA
# as the primary subject). When absent, DO NOT fall back to
# modelVersionIds[0] — that array mixes checkpoints, LoRAs,
# and other model version IDs without ordering guarantees.
# The downstream enrichment flow will find the real
# checkpoint via meta.resources (type:"model" hash) or
# meta.civitaiResources (type:"checkpoint" version ID), so
# leaving model_ver_id as None is safe and avoids the bug
# where a LoRA version ID was treated as the checkpoint.
model_ver_id = image_info.get("modelVersionId")
if not model_ver_id:
ids = image_info.get("modelVersionIds")
if isinstance(ids, list) and ids:
model_ver_id = ids[0]
# Inject root-level modelVersionIds into meta so downstream
# parsers (CivitaiApiMetadataParser) can discover ALL resources