Files
ComfyUI-Lora-Manager/py/recipes/parsers/civitai_image.py
Will Miao 34791c2ad7 fix(recipe): use resources type field to identify checkpoint instead of modelVersionIds[0]
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.
2026-05-28 15:46:38 +08:00

687 lines
30 KiB
Python

"""Parser for Civitai image metadata format."""
import json
import logging
from typing import Dict, Any, Union
from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__)
class CivitaiApiMetadataParser(RecipeMetadataParser):
"""Parser for Civitai image metadata format"""
def is_metadata_matching(self, metadata) -> bool:
"""Check if the metadata matches the Civitai image metadata format
Args:
metadata: The metadata from the image (dict)
Returns:
bool: True if this parser can handle the metadata
"""
if not metadata or not isinstance(metadata, dict):
return False
def has_markers(payload: Dict[str, Any]) -> bool:
# Check for common CivitAI image metadata fields
civitai_image_fields = (
"resources",
"civitaiResources",
"additionalResources",
"hashes",
"prompt",
"negativePrompt",
"steps",
"sampler",
"cfgScale",
"seed",
"width",
"height",
"Model",
"Model hash",
"modelVersionIds",
)
return any(key in payload for key in civitai_image_fields)
# Check the main metadata object
if has_markers(metadata):
return True
# Check for LoRA hash patterns
hashes = metadata.get("hashes")
if isinstance(hashes, dict) and any(
str(key).lower().startswith("lora:") for key in hashes
):
return True
# Check nested meta object (common in CivitAI image responses)
nested_meta = metadata.get("meta")
if isinstance(nested_meta, dict):
if has_markers(nested_meta):
return True
# Also check for LoRA hash patterns in nested meta
hashes = nested_meta.get("hashes")
if isinstance(hashes, dict) and any(
str(key).lower().startswith("lora:") for key in hashes
):
return True
return False
async def parse_metadata( # type: ignore[override]
self, user_comment, recipe_scanner=None, civitai_client=None
) -> Dict[str, Any]:
"""Parse metadata from Civitai image format
Args:
user_comment: The metadata from the image (dict)
recipe_scanner: Optional recipe scanner service
civitai_client: Optional Civitai API client (deprecated, use metadata_provider instead)
Returns:
Dict containing parsed recipe data
"""
metadata: Dict[str, Any] = user_comment # type: ignore[assignment]
metadata = user_comment
try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
# Civitai image responses may wrap the actual metadata inside a "meta" key
if (
isinstance(metadata, dict)
and "meta" in metadata
and isinstance(metadata["meta"], dict)
):
inner_meta = metadata["meta"]
if any(
key in inner_meta
for key in (
"resources",
"civitaiResources",
"additionalResources",
"hashes",
"prompt",
"negativePrompt",
)
):
metadata = inner_meta
# Initialize result structure
result = {
"base_model": None,
"loras": [],
"model": None,
"gen_params": {},
"from_civitai_image": True,
}
# Track already added LoRAs to prevent duplicates
added_loras = {} # key: model_version_id or hash, value: index in result["loras"]
# Extract hash information from hashes field for LoRA matching
lora_hashes = {}
if "hashes" in metadata and isinstance(metadata["hashes"], dict):
for key, hash_value in metadata["hashes"].items():
key_str = str(key)
if key_str.lower().startswith("lora:"):
lora_name = key_str.split(":", 1)[1]
lora_hashes[lora_name] = hash_value
# Extract prompt and negative prompt
if "prompt" in metadata:
result["gen_params"]["prompt"] = metadata["prompt"]
if "negativePrompt" in metadata:
result["gen_params"]["negative_prompt"] = metadata["negativePrompt"]
# Extract other generation parameters
param_mapping = {
"steps": "steps",
"sampler": "sampler",
"cfgScale": "cfg_scale",
"seed": "seed",
"Size": "size",
"clipSkip": "clip_skip",
}
for civitai_key, our_key in param_mapping.items():
if civitai_key in metadata and our_key in GEN_PARAM_KEYS:
result["gen_params"][our_key] = metadata[civitai_key]
# Extract base model information - directly if available
if "baseModel" in metadata:
result["base_model"] = metadata["baseModel"]
elif "Model hash" in metadata and metadata_provider:
model_hash = metadata["Model hash"]
model_info, error = await metadata_provider.get_model_by_hash(
model_hash
)
if model_info:
result["base_model"] = model_info.get("baseModel", "")
elif "Model" in metadata and isinstance(metadata.get("resources"), list):
# Try to find base model in resources
for resource in metadata.get("resources", []):
if resource.get("type") == "model" and resource.get(
"name"
) == metadata.get("Model"):
# This is likely the checkpoint model
if metadata_provider and resource.get("hash"):
(
model_info,
error,
) = await metadata_provider.get_model_by_hash(
resource.get("hash")
)
if model_info:
result["base_model"] = model_info.get("baseModel", "")
base_model_counts = {}
# 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_type == "lora":
lora_hash = resource.get("hash", "")
# Try to get hash from the hashes field if not present in resource
if not lora_hash and resource.get("name"):
lora_hash = lora_hashes.get(resource["name"], "")
# Skip LoRAs without proper identification (hash or modelVersionId)
if not lora_hash and not resource.get("modelVersionId"):
logger.debug(
f"Skipping LoRA resource '{resource.get('name', 'Unknown')}' - no hash or modelVersionId"
)
continue
# Skip if we've already added this LoRA by hash
if lora_hash and lora_hash in added_loras:
continue
lora_entry = {
"name": resource.get("name", "Unknown LoRA"),
"type": "lora",
"weight": float(resource.get("weight", 1.0)),
"hash": lora_hash,
"existsLocally": False,
"localPath": None,
"file_name": resource.get("name", "Unknown"),
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# Try to get info from Civitai if hash is available
if lora_entry["hash"] and metadata_provider:
try:
civitai_info = (
await metadata_provider.get_model_by_hash(lora_hash)
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
continue # Skip invalid LoRA types
lora_entry = populated_entry
# If we have a version ID from Civitai, track it for deduplication
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(
result["loras"]
)
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
)
# Track by hash if we have it
if lora_hash:
added_loras[lora_hash] = len(result["loras"])
result["loras"].append(lora_entry)
# Process civitaiResources array
if "civitaiResources" in metadata and isinstance(
metadata["civitaiResources"], list
):
for resource in metadata["civitaiResources"]:
# Get resource type and identifier
resource_type = str(resource.get("type") or "").lower()
version_id = str(resource.get("modelVersionId", ""))
if resource_type == "checkpoint":
checkpoint_entry = {
"id": resource.get("modelVersionId", 0),
"modelId": resource.get("modelId", 0),
"name": resource.get("modelName", "Unknown Checkpoint"),
"version": resource.get("modelVersionName", ""),
"type": resource.get("type", "checkpoint"),
"existsLocally": False,
"localPath": None,
"file_name": resource.get("modelName", ""),
"hash": resource.get("hash", "") or "",
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
if version_id and metadata_provider:
try:
civitai_info = (
await metadata_provider.get_model_version_info(
version_id
)
)
checkpoint_entry = (
await self.populate_checkpoint_from_civitai(
checkpoint_entry, civitai_info
)
)
except Exception as e:
logger.error(
f"Error fetching Civitai info for checkpoint version {version_id}: {e}"
)
if result["model"] is None:
result["model"] = checkpoint_entry
continue
# Skip if we've already added this LoRA
if version_id and version_id in added_loras:
continue
# Initialize lora entry
lora_entry = {
"id": resource.get("modelVersionId", 0),
"modelId": resource.get("modelId", 0),
"name": resource.get("modelName", "Unknown LoRA"),
"version": resource.get("modelVersionName", ""),
"type": resource.get("type", "lora"),
"weight": round(float(resource.get("weight", 1.0)), 2),
"existsLocally": False,
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# Try to get info from Civitai if modelVersionId is available
if version_id and metadata_provider:
try:
# Use get_model_version_info instead of get_model_version
civitai_info = (
await metadata_provider.get_model_version_info(
version_id
)
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
)
if populated_entry is None:
continue # Skip invalid LoRA types
lora_entry = populated_entry
except Exception as e:
logger.error(
f"Error fetching Civitai info for model version {version_id}: {e}"
)
# Track this LoRA in our deduplication dict
if version_id:
added_loras[version_id] = len(result["loras"])
result["loras"].append(lora_entry)
# Process additionalResources array
if "additionalResources" in metadata and isinstance(
metadata["additionalResources"], list
):
for resource in metadata["additionalResources"]:
# Skip resources that aren't LoRAs or LyCORIS
if (
resource.get("type") not in ["lora", "lycoris"]
and "type" not in resource
):
continue
lora_type = resource.get("type", "lora")
name = resource.get("name", "")
# Extract ID from URN format if available
version_id = None
if name and "civitai:" in name:
parts = name.split("@")
if len(parts) > 1:
version_id = parts[1]
# Skip if we've already added this LoRA
if version_id in added_loras:
continue
lora_entry = {
"name": name,
"type": lora_type,
"weight": float(resource.get("strength", 1.0)),
"hash": "",
"existsLocally": False,
"localPath": None,
"file_name": name,
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# If we have a version ID and metadata provider, try to get more info
if version_id and metadata_provider:
try:
# Use get_model_version_info with the version ID
civitai_info = (
await metadata_provider.get_model_version_info(
version_id
)
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
)
if populated_entry is None:
continue # Skip invalid LoRA types
lora_entry = populated_entry
# Track this LoRA for deduplication
if version_id:
added_loras[version_id] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for model ID {version_id}: {e}"
)
result["loras"].append(lora_entry)
# Process modelVersionIds from Civitai image API
# These are model version IDs returned at root level when meta doesn't contain resources
if "modelVersionIds" in metadata and isinstance(
metadata["modelVersionIds"], list
):
for version_id in metadata["modelVersionIds"]:
version_id_str = str(version_id)
# Skip if we've already added this LoRA by version ID
if version_id_str in added_loras:
continue
# Initialize lora entry with version ID
lora_entry = {
"id": version_id,
"modelId": 0,
"name": "Unknown LoRA",
"version": "",
"type": "lora",
"weight": 1.0,
"existsLocally": False,
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# Fetch model info from Civitai
if metadata_provider and version_id_str:
try:
civitai_info = (
await metadata_provider.get_model_version_info(
version_id_str
)
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
)
if populated_entry is None:
continue # Skip invalid LoRA types
lora_entry = populated_entry
except Exception as e:
logger.error(
f"Error fetching Civitai info for model version {version_id}: {e}"
)
# Track this LoRA for deduplication
if version_id_str:
added_loras[version_id_str] = len(result["loras"])
result["loras"].append(lora_entry)
# If we found LoRA hashes in the metadata but haven't already
# populated entries for them, fall back to creating LoRAs from
# the hashes section. Some Civitai image responses only include
# LoRA information here without explicit resources entries.
for lora_name, lora_hash in lora_hashes.items():
if not lora_hash:
continue
# Skip LoRAs we've already added via resources or other fields
if lora_hash in added_loras:
continue
lora_entry = {
"name": lora_name,
"type": "lora",
"weight": 1.0,
"hash": lora_hash,
"existsLocally": False,
"localPath": None,
"file_name": lora_name,
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
if metadata_provider:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
continue
lora_entry = populated_entry
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_hash}: {e}"
)
added_loras[lora_hash] = len(result["loras"])
result["loras"].append(lora_entry)
# Check for LoRA info in the format "Lora_0 Model hash", "Lora_0 Model name", etc.
lora_index = 0
while (
f"Lora_{lora_index} Model hash" in metadata
and f"Lora_{lora_index} Model name" in metadata
):
lora_hash = metadata[f"Lora_{lora_index} Model hash"]
lora_name = metadata[f"Lora_{lora_index} Model name"]
lora_strength_model = float(
metadata.get(f"Lora_{lora_index} Strength model", 1.0)
)
# Skip if we've already added this LoRA by hash
if lora_hash and lora_hash in added_loras:
lora_index += 1
continue
lora_entry = {
"name": lora_name,
"type": "lora",
"weight": lora_strength_model,
"hash": lora_hash,
"existsLocally": False,
"localPath": None,
"file_name": lora_name,
"thumbnailUrl": "/loras_static/images/no-preview.png",
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
# Try to get info from Civitai if hash is available
if lora_entry["hash"] and metadata_provider:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
lora_index += 1
continue # Skip invalid LoRA types
lora_entry = populated_entry
# If we have a version ID from Civitai, track it for deduplication
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
)
# Track by hash if we have it
if lora_hash:
added_loras[lora_hash] = len(result["loras"])
result["loras"].append(lora_entry)
lora_index += 1
# If base model wasn't found earlier, use the most common one from LoRAs
if not result["base_model"] and base_model_counts:
result["base_model"] = max(
base_model_counts.items(), key=lambda x: x[1]
)[0]
return result
except Exception as e:
logger.error(f"Error parsing Civitai image metadata: {e}", exc_info=True)
return {"error": str(e), "loras": []}