From 3338c17e8f289aa96904e6b85a245395228de97a Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Sun, 16 Mar 2025 18:56:27 +0800 Subject: [PATCH] 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. --- py/routes/recipe_routes.py | 17 ++++++--- py/services/recipe_scanner.py | 67 ----------------------------------- 2 files changed, 12 insertions(+), 72 deletions(-) diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py index 8d6b1369..177b9554 100644 --- a/py/routes/recipe_routes.py +++ b/py/routes/recipe_routes.py @@ -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', ''), diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index 728c2038..88421e7b 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -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"""