mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 15:15:44 -03:00
feat: Enhance Civitai image metadata parser to prevent duplicate LoRAs
This commit is contained in:
@@ -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)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user