mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
Refactor recipe processing in RecipeRoutes to enhance LoRA handling. Introduce base model counting logic to determine the most common base model from LoRAs, and streamline the collection of LoRA metadata. Remove outdated metadata update method from RecipeScanner to improve code clarity and maintainability.
This commit is contained in:
@@ -217,9 +217,9 @@ class RecipeRoutes:
|
||||
loras = []
|
||||
base_model = None
|
||||
|
||||
# Set base model from checkpoint if available
|
||||
if checkpoint:
|
||||
base_model = checkpoint.get('modelName', '')
|
||||
# Process LoRAs and collect base models
|
||||
base_model_counts = {}
|
||||
loras = []
|
||||
|
||||
# Process LoRAs
|
||||
for resource in civitai_resources:
|
||||
@@ -278,8 +278,11 @@ class RecipeRoutes:
|
||||
if 'images' in civitai_info and civitai_info['images']:
|
||||
lora_entry['thumbnailUrl'] = civitai_info['images'][0].get('url', '')
|
||||
|
||||
# Get base model
|
||||
lora_entry['baseModel'] = civitai_info.get('baseModel', '')
|
||||
# Get base model and update counts
|
||||
current_base_model = civitai_info.get('baseModel', '')
|
||||
lora_entry['baseModel'] = current_base_model
|
||||
if current_base_model:
|
||||
base_model_counts[current_base_model] = base_model_counts.get(current_base_model, 0) + 1
|
||||
|
||||
# Get download URL
|
||||
lora_entry['downloadUrl'] = civitai_info.get('downloadUrl', '')
|
||||
@@ -290,6 +293,10 @@ class RecipeRoutes:
|
||||
|
||||
loras.append(lora_entry)
|
||||
|
||||
# Set base_model to the most common one from civitai_info
|
||||
if base_model_counts:
|
||||
base_model = max(base_model_counts.items(), key=lambda x: x[1])[0]
|
||||
|
||||
# Extract generation parameters for recipe metadata
|
||||
gen_params = {
|
||||
'prompt': metadata.get('prompt', ''),
|
||||
|
||||
@@ -317,73 +317,6 @@ class RecipeScanner:
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting hash from Civitai: {e}")
|
||||
return None
|
||||
|
||||
async def _update_recipe_metadata(self, recipe_data: Dict) -> bool:
|
||||
"""Update recipe metadata with missing information
|
||||
|
||||
Returns:
|
||||
bool: True if metadata was updated
|
||||
"""
|
||||
metadata_updated = False
|
||||
|
||||
# Update lora information
|
||||
for lora in recipe_data.get('loras', []):
|
||||
# First check if modelVersionId exists and hash doesn't
|
||||
if 'modelVersionId' in lora and not lora.get('hash'):
|
||||
model_version_id = str(lora['modelVersionId'])
|
||||
# Try to find hash in lora cache first
|
||||
hash_from_cache = await self._find_hash_in_lora_cache(model_version_id)
|
||||
if hash_from_cache:
|
||||
logger.info(f"Found hash in cache for modelVersionId {model_version_id}")
|
||||
lora['hash'] = hash_from_cache.lower() # Standardize to lowercase
|
||||
metadata_updated = True
|
||||
else:
|
||||
# If not in cache, fetch from Civitai
|
||||
logger.info(f"Fetching hash from Civitai for {model_version_id}")
|
||||
hash_from_civitai = await self._get_hash_from_civitai(model_version_id)
|
||||
if hash_from_civitai:
|
||||
logger.info(f"Got hash from Civitai")
|
||||
lora['hash'] = hash_from_civitai.lower() # Standardize to lowercase
|
||||
metadata_updated = True
|
||||
else:
|
||||
logger.warning(f"Could not get hash for modelVersionId {model_version_id}")
|
||||
|
||||
# If modelVersionId exists but no modelVersionName, try to get it from Civitai
|
||||
if 'modelVersionId' in lora and not lora.get('modelVersionName'):
|
||||
model_version_id = str(lora['modelVersionId'])
|
||||
model_version_name = await self._get_model_version_name(model_version_id)
|
||||
if model_version_name:
|
||||
lora['modelVersionName'] = model_version_name
|
||||
metadata_updated = True
|
||||
|
||||
# If has hash, check if it's in library
|
||||
if 'hash' in lora:
|
||||
hash_value = lora['hash'].lower() # Ensure lowercase when comparing
|
||||
in_library = self._lora_scanner.has_lora_hash(hash_value)
|
||||
lora['inLibrary'] = in_library
|
||||
|
||||
# If hash is in library but no file_name, look up and set file_name
|
||||
if in_library and (not lora.get('file_name') or not lora['file_name']):
|
||||
lora_path = self._lora_scanner.get_lora_path_by_hash(hash_value)
|
||||
if lora_path:
|
||||
file_name = os.path.splitext(os.path.basename(lora_path))[0]
|
||||
logger.info(f"Found lora in library: {file_name}")
|
||||
lora['file_name'] = file_name
|
||||
metadata_updated = True
|
||||
elif not in_library:
|
||||
# Lora not in library
|
||||
logger.info(f"LoRA with hash {hash_value[:8]}... not found in library")
|
||||
lora['file_name'] = ''
|
||||
metadata_updated = True
|
||||
|
||||
# Determine the base_model for the recipe based on loras
|
||||
if recipe_data.get('loras') and not recipe_data.get('base_model'):
|
||||
base_model = await self._determine_base_model(recipe_data.get('loras', []))
|
||||
if base_model:
|
||||
recipe_data['base_model'] = base_model
|
||||
metadata_updated = True
|
||||
|
||||
return metadata_updated
|
||||
|
||||
async def _get_model_version_name(self, model_version_id: str) -> Optional[str]:
|
||||
"""Get model version name from Civitai API"""
|
||||
|
||||
Reference in New Issue
Block a user