feat: Enhance Civitai image metadata parser to prevent duplicate LoRAs

This commit is contained in:
Will Miao
2025-07-02 16:50:19 +08:00
parent 2d4f6ae7ce
commit 9d8b7344cd

View File

@@ -50,6 +50,9 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
'from_civitai_image': True '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 prompt and negative prompt # Extract prompt and negative prompt
if "prompt" in metadata: if "prompt" in metadata:
result["gen_params"]["prompt"] = metadata["prompt"] result["gen_params"]["prompt"] = metadata["prompt"]
@@ -96,11 +99,17 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
for resource in metadata["resources"]: for resource in metadata["resources"]:
# Modified to process resources without a type field as potential LoRAs # Modified to process resources without a type field as potential LoRAs
if resource.get("type", "lora") == "lora": if resource.get("type", "lora") == "lora":
lora_hash = resource.get("hash", "")
# Skip if we've already added this LoRA by hash
if lora_hash and lora_hash in added_loras:
continue
lora_entry = { lora_entry = {
'name': resource.get("name", "Unknown LoRA"), 'name': resource.get("name", "Unknown LoRA"),
'type': "lora", 'type': "lora",
'weight': float(resource.get("weight", 1.0)), 'weight': float(resource.get("weight", 1.0)),
'hash': resource.get("hash", ""), 'hash': lora_hash,
'existsLocally': False, 'existsLocally': False,
'localPath': None, 'localPath': None,
'file_name': resource.get("name", "Unknown"), 'file_name': resource.get("name", "Unknown"),
@@ -114,7 +123,6 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Try to get info from Civitai if hash is available # Try to get info from Civitai if hash is available
if lora_entry['hash'] and civitai_client: if lora_entry['hash'] and civitai_client:
try: try:
lora_hash = lora_entry['hash']
civitai_info = await civitai_client.get_model_by_hash(lora_hash) civitai_info = await civitai_client.get_model_by_hash(lora_hash)
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
@@ -129,17 +137,34 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
continue # Skip invalid LoRA types continue # Skip invalid LoRA types
lora_entry = populated_entry 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: except Exception as e:
logger.error(f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {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) result["loras"].append(lora_entry)
# Process civitaiResources array # Process civitaiResources array
if "civitaiResources" in metadata and isinstance(metadata["civitaiResources"], list): if "civitaiResources" in metadata and isinstance(metadata["civitaiResources"], list):
for resource in metadata["civitaiResources"]: for resource in metadata["civitaiResources"]:
# Modified to process resources without a type field as potential LoRAs # Skip resources that aren't LoRAs or LyCORIS
if resource.get("type") in ["lora", "lycoris"] or "type" not in resource: if resource.get("type") not in ["lora", "lycoris"] and "type" not in resource:
# Initialize lora entry with the same structure as in automatic.py continue
# Get unique identifier for deduplication
version_id = str(resource.get("modelVersionId", ""))
# Skip if we've already added this LoRA
if version_id and version_id in added_loras:
continue
# Initialize lora entry
lora_entry = { lora_entry = {
'id': resource.get("modelVersionId", 0), 'id': resource.get("modelVersionId", 0),
'modelId': resource.get("modelId", 0), 'modelId': resource.get("modelId", 0),
@@ -156,9 +181,8 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
} }
# Try to get info from Civitai if modelVersionId is available # Try to get info from Civitai if modelVersionId is available
if resource.get('modelVersionId') and civitai_client: if version_id and civitai_client:
try: try:
version_id = str(resource.get('modelVersionId'))
# Use get_model_version_info instead of get_model_version # Use get_model_version_info instead of get_model_version
civitai_info, error = await civitai_client.get_model_version_info(version_id) civitai_info, error = await civitai_client.get_model_version_info(version_id)
@@ -178,24 +202,34 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
lora_entry = populated_entry lora_entry = populated_entry
except Exception as e: except Exception as e:
logger.error(f"Error fetching Civitai info for model version {resource.get('modelVersionId')}: {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) result["loras"].append(lora_entry)
# Process additionalResources array # Process additionalResources array
if "additionalResources" in metadata and isinstance(metadata["additionalResources"], list): if "additionalResources" in metadata and isinstance(metadata["additionalResources"], list):
for resource in metadata["additionalResources"]: for resource in metadata["additionalResources"]:
# Modified to process resources without a type field as potential LoRAs # Skip resources that aren't LoRAs or LyCORIS
if resource.get("type") in ["lora", "lycoris"] or "type" not in resource: if resource.get("type") not in ["lora", "lycoris"] and "type" not in resource:
continue
lora_type = resource.get("type", "lora") lora_type = resource.get("type", "lora")
name = resource.get("name", "") name = resource.get("name", "")
# Extract ID from URN format if available # Extract ID from URN format if available
model_id = None version_id = None
if name and "civitai:" in name: if name and "civitai:" in name:
parts = name.split("@") parts = name.split("@")
if len(parts) > 1: if len(parts) > 1:
model_id = parts[1] version_id = parts[1]
# Skip if we've already added this LoRA
if version_id in added_loras:
continue
lora_entry = { lora_entry = {
'name': name, 'name': name,
@@ -212,11 +246,11 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
'isDeleted': False 'isDeleted': False
} }
# If we have a model ID and civitai client, try to get more info # If we have a version ID and civitai client, try to get more info
if model_id and civitai_client: if version_id and civitai_client:
try: try:
# Use get_model_version_info with the model ID # Use get_model_version_info with the version ID
civitai_info, error = await civitai_client.get_model_version_info(model_id) civitai_info, error = await civitai_client.get_model_version_info(version_id)
if error: if error:
logger.warning(f"Error getting model version info: {error}") logger.warning(f"Error getting model version info: {error}")
@@ -232,8 +266,12 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
continue # Skip invalid LoRA types continue # Skip invalid LoRA types
lora_entry = populated_entry lora_entry = populated_entry
# Track this LoRA for deduplication
if version_id:
added_loras[version_id] = len(result["loras"])
except Exception as e: except Exception as e:
logger.error(f"Error fetching Civitai info for model ID {model_id}: {e}") logger.error(f"Error fetching Civitai info for model ID {version_id}: {e}")
result["loras"].append(lora_entry) result["loras"].append(lora_entry)