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:
Will Miao
2025-03-16 18:56:27 +08:00
parent 22085e5174
commit 3338c17e8f
2 changed files with 12 additions and 72 deletions

View File

@@ -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', ''),

View File

@@ -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"""