mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-24 06:32:12 -03:00
feat(parser): enhance model metadata extraction in Automatic1111 parser
- Add MODEL_NAME_PATTERN regex to extract model names from parameters - Extract model hash from parsed hashes when available in metadata - Add checkpoint model hash and name extraction from parameters section - Implement checkpoint resource processing from Civitai metadata - Improve model information completeness for better recipe tracking
This commit is contained in:
Will Miao
@@ -1,6 +1,7 @@
|
||||
"""Parser for Automatic1111 metadata format."""
|
||||
|
||||
import re
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, Any
|
||||
@@ -22,6 +23,7 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
CIVITAI_METADATA_REGEX = r', Civitai metadata:\s*(\{.*?\})'
|
||||
EXTRANETS_REGEX = r'<(lora|hypernet):([^:]+):(-?[0-9.]+)>'
|
||||
MODEL_HASH_PATTERN = r'Model hash: ([a-zA-Z0-9]+)'
|
||||
MODEL_NAME_PATTERN = r'Model: ([^,]+)'
|
||||
VAE_HASH_PATTERN = r'VAE hash: ([a-zA-Z0-9]+)'
|
||||
|
||||
def is_metadata_matching(self, user_comment: str) -> bool:
|
||||
@@ -115,6 +117,12 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
except json.JSONDecodeError:
|
||||
logger.error("Error parsing hashes JSON")
|
||||
|
||||
# Pick up model hash from parsed hashes if available
|
||||
if "hashes" in metadata and not metadata.get("model_hash"):
|
||||
model_hash_from_hashes = metadata["hashes"].get("model")
|
||||
if model_hash_from_hashes:
|
||||
metadata["model_hash"] = model_hash_from_hashes
|
||||
|
||||
# Extract Lora hashes in alternative format
|
||||
lora_hashes_match = re.search(self.LORA_HASHES_REGEX, params_section)
|
||||
if not hashes_match and lora_hashes_match:
|
||||
@@ -137,6 +145,17 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
params_section = params_section.replace(lora_hashes_match.group(0), '')
|
||||
except Exception as e:
|
||||
logger.error(f"Error parsing Lora hashes: {e}")
|
||||
|
||||
# Extract checkpoint model hash/name when provided outside Civitai resources
|
||||
model_hash_match = re.search(self.MODEL_HASH_PATTERN, params_section)
|
||||
if model_hash_match:
|
||||
metadata["model_hash"] = model_hash_match.group(1).strip()
|
||||
params_section = params_section.replace(model_hash_match.group(0), '')
|
||||
|
||||
model_name_match = re.search(self.MODEL_NAME_PATTERN, params_section)
|
||||
if model_name_match:
|
||||
metadata["model_name"] = model_name_match.group(1).strip()
|
||||
params_section = params_section.replace(model_name_match.group(0), '')
|
||||
|
||||
# Extract basic parameters
|
||||
param_pattern = r'([A-Za-z\s]+): ([^,]+)'
|
||||
@@ -178,9 +197,10 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
|
||||
metadata["gen_params"] = gen_params
|
||||
|
||||
# Extract LoRA information
|
||||
# Extract LoRA and checkpoint information
|
||||
loras = []
|
||||
base_model_counts = {}
|
||||
checkpoint = None
|
||||
|
||||
# First use Civitai resources if available (more reliable source)
|
||||
if metadata.get("civitai_resources"):
|
||||
@@ -202,6 +222,50 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
resource["modelVersionId"] = air_modelVersionId
|
||||
# --- End added ---
|
||||
|
||||
if resource.get("type") == "checkpoint" and resource.get("modelVersionId"):
|
||||
version_id = resource.get("modelVersionId")
|
||||
version_id_str = str(version_id)
|
||||
checkpoint_entry = {
|
||||
'id': version_id,
|
||||
'modelId': resource.get("modelId", 0),
|
||||
'name': resource.get("modelName", "Unknown Checkpoint"),
|
||||
'version': resource.get("modelVersionName", resource.get("versionName", "")),
|
||||
'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 metadata_provider:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_version_info(version_id_str)
|
||||
checkpoint_entry = await self.populate_checkpoint_from_civitai(
|
||||
checkpoint_entry,
|
||||
civitai_info
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error fetching Civitai info for checkpoint version %s: %s",
|
||||
version_id,
|
||||
e,
|
||||
)
|
||||
|
||||
# Prefer the first checkpoint found
|
||||
if checkpoint_entry.get("baseModel"):
|
||||
base_model_value = checkpoint_entry["baseModel"]
|
||||
base_model_counts[base_model_value] = base_model_counts.get(base_model_value, 0) + 1
|
||||
|
||||
if checkpoint is None:
|
||||
checkpoint = checkpoint_entry
|
||||
|
||||
continue
|
||||
|
||||
if resource.get("type") in ["lora", "lycoris", "hypernet"] and resource.get("modelVersionId"):
|
||||
# Initialize lora entry
|
||||
lora_entry = {
|
||||
@@ -237,6 +301,52 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
|
||||
loras.append(lora_entry)
|
||||
|
||||
# Fallback checkpoint parsing from generic "Model" and "Model hash" fields
|
||||
if checkpoint is None:
|
||||
model_hash = metadata.get("model_hash")
|
||||
if not model_hash and metadata.get("hashes"):
|
||||
model_hash = metadata["hashes"].get("model")
|
||||
|
||||
model_name = metadata.get("model_name")
|
||||
file_name = ""
|
||||
if model_name:
|
||||
cleaned_name = re.split(r"[\\\\/]", model_name)[-1]
|
||||
file_name = os.path.splitext(cleaned_name)[0]
|
||||
|
||||
if model_hash or model_name:
|
||||
checkpoint_entry = {
|
||||
'id': 0,
|
||||
'modelId': 0,
|
||||
'name': model_name or "Unknown Checkpoint",
|
||||
'version': '',
|
||||
'type': 'checkpoint',
|
||||
'hash': model_hash or "",
|
||||
'existsLocally': False,
|
||||
'localPath': None,
|
||||
'file_name': file_name,
|
||||
'thumbnailUrl': '/loras_static/images/no-preview.png',
|
||||
'baseModel': '',
|
||||
'size': 0,
|
||||
'downloadUrl': '',
|
||||
'isDeleted': False
|
||||
}
|
||||
|
||||
if metadata_provider and model_hash:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(model_hash)
|
||||
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 hash {model_hash}: {e}")
|
||||
|
||||
if checkpoint_entry.get("baseModel"):
|
||||
base_model_value = checkpoint_entry["baseModel"]
|
||||
base_model_counts[base_model_value] = base_model_counts.get(base_model_value, 0) + 1
|
||||
|
||||
checkpoint = checkpoint_entry
|
||||
|
||||
# If no LoRAs from Civitai resources or to supplement, extract from metadata["hashes"]
|
||||
if not loras or len(loras) == 0:
|
||||
# Extract lora weights from extranet tags in prompt (for later use)
|
||||
@@ -300,7 +410,9 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
|
||||
# Try to get base model from resources or make educated guess
|
||||
base_model = None
|
||||
if base_model_counts:
|
||||
if checkpoint and checkpoint.get("baseModel"):
|
||||
base_model = checkpoint.get("baseModel")
|
||||
elif base_model_counts:
|
||||
# Use the most common base model from the loras
|
||||
base_model = max(base_model_counts.items(), key=lambda x: x[1])[0]
|
||||
|
||||
@@ -317,6 +429,10 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
'gen_params': filtered_gen_params,
|
||||
'from_automatic_metadata': True
|
||||
}
|
||||
|
||||
if checkpoint:
|
||||
result['checkpoint'] = checkpoint
|
||||
result['model'] = checkpoint
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Parser for meta format (Lora_N Model hash) metadata."""
|
||||
|
||||
import os
|
||||
import re
|
||||
import logging
|
||||
from typing import Dict, Any
|
||||
@@ -145,14 +146,53 @@ class MetaFormatParser(RecipeMetadataParser):
|
||||
|
||||
loras.append(lora_entry)
|
||||
|
||||
# Extract model information
|
||||
model = None
|
||||
if 'model' in metadata:
|
||||
model = metadata['model']
|
||||
# Extract checkpoint information from generic Model/Model hash fields
|
||||
checkpoint = None
|
||||
model_hash = metadata.get("model_hash")
|
||||
model_name = metadata.get("model")
|
||||
|
||||
if model_hash or model_name:
|
||||
cleaned_name = None
|
||||
if model_name:
|
||||
cleaned_name = re.split(r"[\\\\/]", model_name)[-1]
|
||||
cleaned_name = os.path.splitext(cleaned_name)[0]
|
||||
|
||||
checkpoint_entry = {
|
||||
'id': 0,
|
||||
'modelId': 0,
|
||||
'name': model_name or "Unknown Checkpoint",
|
||||
'version': '',
|
||||
'type': 'checkpoint',
|
||||
'hash': model_hash or "",
|
||||
'existsLocally': False,
|
||||
'localPath': None,
|
||||
'file_name': cleaned_name or (model_name or ""),
|
||||
'thumbnailUrl': '/loras_static/images/no-preview.png',
|
||||
'baseModel': '',
|
||||
'size': 0,
|
||||
'downloadUrl': '',
|
||||
'isDeleted': False
|
||||
}
|
||||
|
||||
if metadata_provider and model_hash:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(model_hash)
|
||||
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 hash {model_hash}: {e}")
|
||||
|
||||
if checkpoint_entry.get("baseModel"):
|
||||
base_model_value = checkpoint_entry["baseModel"]
|
||||
base_model_counts[base_model_value] = base_model_counts.get(base_model_value, 0) + 1
|
||||
|
||||
checkpoint = checkpoint_entry
|
||||
|
||||
# Set base_model to the most common one from civitai_info
|
||||
base_model = None
|
||||
if base_model_counts:
|
||||
# Set base_model to the most common one from civitai_info or checkpoint
|
||||
base_model = checkpoint["baseModel"] if checkpoint and checkpoint.get("baseModel") else None
|
||||
if not base_model and base_model_counts:
|
||||
base_model = max(base_model_counts.items(), key=lambda x: x[1])[0]
|
||||
|
||||
# Extract generation parameters for recipe metadata
|
||||
@@ -170,7 +210,8 @@ class MetaFormatParser(RecipeMetadataParser):
|
||||
'loras': loras,
|
||||
'gen_params': gen_params,
|
||||
'raw_metadata': metadata,
|
||||
'from_meta_format': True
|
||||
'from_meta_format': True,
|
||||
**({'checkpoint': checkpoint, 'model': checkpoint} if checkpoint else {})
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
|
||||
@@ -94,8 +94,45 @@ class RecipeFormatParser(RecipeMetadataParser):
|
||||
lora_entry['thumbnailUrl'] = '/loras_static/images/no-preview.png'
|
||||
|
||||
loras.append(lora_entry)
|
||||
|
||||
|
||||
logger.info(f"Found {len(loras)} loras in recipe metadata")
|
||||
|
||||
# Process checkpoint information if present
|
||||
checkpoint = None
|
||||
checkpoint_data = recipe_metadata.get('checkpoint') or {}
|
||||
if isinstance(checkpoint_data, dict) and checkpoint_data:
|
||||
version_id = checkpoint_data.get('modelVersionId') or checkpoint_data.get('id')
|
||||
checkpoint_entry = {
|
||||
'id': version_id or 0,
|
||||
'modelId': checkpoint_data.get('modelId', 0),
|
||||
'name': checkpoint_data.get('name', 'Unknown Checkpoint'),
|
||||
'version': checkpoint_data.get('version', ''),
|
||||
'type': checkpoint_data.get('type', 'checkpoint'),
|
||||
'hash': checkpoint_data.get('hash', ''),
|
||||
'existsLocally': False,
|
||||
'localPath': None,
|
||||
'file_name': checkpoint_data.get('file_name', ''),
|
||||
'thumbnailUrl': '/loras_static/images/no-preview.png',
|
||||
'baseModel': '',
|
||||
'size': 0,
|
||||
'downloadUrl': '',
|
||||
'isDeleted': False
|
||||
}
|
||||
|
||||
if metadata_provider:
|
||||
try:
|
||||
civitai_info = None
|
||||
if version_id:
|
||||
civitai_info = await metadata_provider.get_model_version_info(str(version_id))
|
||||
elif checkpoint_entry.get('hash'):
|
||||
civitai_info = await metadata_provider.get_model_by_hash(checkpoint_entry['hash'])
|
||||
|
||||
if civitai_info:
|
||||
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 in recipe metadata: {e}")
|
||||
|
||||
checkpoint = checkpoint_entry
|
||||
|
||||
# Filter gen_params to only include recognized keys
|
||||
filtered_gen_params = {}
|
||||
@@ -105,12 +142,13 @@ class RecipeFormatParser(RecipeMetadataParser):
|
||||
filtered_gen_params[key] = value
|
||||
|
||||
return {
|
||||
'base_model': recipe_metadata.get('base_model', ''),
|
||||
'base_model': checkpoint['baseModel'] if checkpoint and checkpoint.get('baseModel') else recipe_metadata.get('base_model', ''),
|
||||
'loras': loras,
|
||||
'gen_params': filtered_gen_params,
|
||||
'tags': recipe_metadata.get('tags', []),
|
||||
'title': recipe_metadata.get('title', ''),
|
||||
'from_recipe_metadata': True
|
||||
'from_recipe_metadata': True,
|
||||
**({'checkpoint': checkpoint, 'model': checkpoint} if checkpoint else {})
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
|
||||
Reference in New Issue
Block a user