From b0c7a1baae57b773f88e69a79d51931266dd9c22 Mon Sep 17 00:00:00 2001 From: Aaalice <65744111+Aaalice233@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:07:17 +0800 Subject: [PATCH] Fix recipe parsing for metadata-free local LoRAs (#1065) * fix(recipes): resolve metadata-free local LoRAs * fix(recipes): prioritize LoRA hashes over names --- py/recipes/base.py | 34 +++ py/recipes/parsers/automatic.py | 262 ++++++++++++++---- py/recipes/parsers/comfy.py | 163 ++++++----- py/services/recipe_scanner.py | 63 ++++- .../test_automatic_metadata_parser.py | 246 ++++++++++++++++ tests/services/test_comfy_metadata_parser.py | 166 +++++++++++ tests/services/test_recipe_scanner.py | 29 ++ 7 files changed, 831 insertions(+), 132 deletions(-) diff --git a/py/recipes/base.py b/py/recipes/base.py index 3e2364b4..3fb92c42 100644 --- a/py/recipes/base.py +++ b/py/recipes/base.py @@ -41,6 +41,40 @@ class RecipeMetadataParser(ABC): """ pass + @staticmethod + def populate_lora_from_local(lora_entry: Dict[str, Any], local_lora: Dict[str, Any], base_model_counts=None) -> Dict[str, Any]: + """Populate a recipe LoRA entry from the local scanner cache.""" + local_path = local_lora.get('file_path') or '' + file_name = local_lora.get('file_name') or os.path.splitext(os.path.basename(local_path))[0] + base_model = local_lora.get('base_model') or '' + + lora_entry['name'] = local_lora.get('model_name') or file_name or lora_entry.get('name', '') + lora_entry['file_name'] = file_name + lora_entry['hash'] = (local_lora.get('sha256') or lora_entry.get('hash') or '').lower() + lora_entry['localPath'] = local_path or None + lora_entry['size'] = local_lora.get('size', 0) or 0 + lora_entry['baseModel'] = base_model + lora_entry['existsLocally'] = True + lora_entry['isDeleted'] = False + + preview_url = local_lora.get('preview_url') + if preview_url: + lora_entry['thumbnailUrl'] = config.get_preview_static_url(preview_url) + + civitai_info = local_lora.get('civitai') or {} + if isinstance(civitai_info, dict): + if civitai_info.get('id') is not None: + lora_entry['id'] = civitai_info['id'] + if civitai_info.get('modelId') is not None: + lora_entry['modelId'] = civitai_info['modelId'] + if civitai_info.get('name'): + lora_entry['version'] = civitai_info['name'] + + if base_model_counts is not None and base_model: + base_model_counts[base_model] = base_model_counts.get(base_model, 0) + 1 + + return lora_entry + @staticmethod async def populate_lora_from_civitai(lora_entry: Dict[str, Any], civitai_info_tuple: Tuple[Dict[str, Any] | None, str | None] | Dict[str, Any], recipe_scanner=None, base_model_counts=None, hash_value=None) -> Optional[Dict[str, Any]]: diff --git a/py/recipes/parsers/automatic.py b/py/recipes/parsers/automatic.py index 4029f2d5..02b5033d 100644 --- a/py/recipes/parsers/automatic.py +++ b/py/recipes/parsers/automatic.py @@ -362,68 +362,208 @@ class AutomaticMetadataParser(RecipeMetadataParser): 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) - lora_weights = {} - lora_matches = re.findall(self.EXTRANETS_REGEX, prompt) - for lora_type, lora_name, lora_weight in lora_matches: - key = f"{lora_type}:{lora_name}" - lora_weights[key] = round(float(lora_weight), 2) - - # Use hashes from metadata as the primary source - if metadata.get("hashes"): - for hash_key, lora_hash in metadata.get("hashes", {}).items(): - # Only process lora or hypernet types - if not hash_key.startswith(("lora:", "hypernet:")): + def normalize_lora_name(name, basename=False): + normalized = str(name or '').replace('\\', '/') + if normalized.casefold().endswith('.safetensors'): + normalized = normalized[:-12] + if basename: + normalized = normalized.rsplit('/', 1)[-1] + return normalized.casefold() + + def get_version_id(lora): + version_id = lora.get('id') + if version_id in (None, '', 0, '0'): + version_id = lora.get('modelVersionId') + if version_id in (None, '', 0, '0'): + return None + return str(version_id) + + prompt_loras = {} + for match in re.findall(self.EXTRANETS_REGEX, prompt): + lora_type, lora_name, _ = match + prompt_loras[(lora_type, normalize_lora_name(lora_name))] = match + + prompt_by_basename = {} + for lora_type, lora_name, lora_weight in prompt_loras.values(): + key = (lora_type, normalize_lora_name(lora_name, True)) + prompt_by_basename.setdefault(key, []).append((lora_name, round(float(lora_weight), 2))) + + hash_basenames = { + (hash_key.split(':', 1)[0], normalize_lora_name(hash_key.split(':', 1)[1], True)) + for hash_key, hash_value in metadata.get("hashes", {}).items() + if hash_value and hash_key.startswith(("lora:", "hypernet:")) + } + recipe_base_model = checkpoint.get("baseModel") if checkpoint else None + if not recipe_base_model and len(base_model_counts) == 1: + recipe_base_model = next(iter(base_model_counts)) + + resource_lora_count = len(loras) + + def make_lora_entry(lora_type, lora_name, weight, lora_hash=''): + return { + 'name': lora_name, + 'type': lora_type, + 'weight': weight, + 'hash': lora_hash, + 'existsLocally': False, + 'localPath': None, + 'file_name': lora_name, + 'thumbnailUrl': '/loras_static/images/no-preview.png', + 'baseModel': '', + 'size': 0, + 'downloadUrl': '', + 'isDeleted': False + } + + def merge_or_append_civitai(civitai_entry, preserve_existing_weight=False): + civitai_id = get_version_id(civitai_entry) + civitai_hash = (civitai_entry.get('hash') or '').lower() + for index, existing in enumerate(loras): + existing_id = get_version_id(existing) + existing_hash = (existing.get('hash') or '').lower() + if not ( + (civitai_id and existing_id == civitai_id) + or (civitai_hash and existing_hash == civitai_hash) + ): + continue + + if preserve_existing_weight: + civitai_entry['weight'] = existing.get('weight', civitai_entry['weight']) + existing_base = existing.get('baseModel') + if not civitai_entry.get('baseModel'): + civitai_entry['baseModel'] = existing_base or '' + elif existing_base: + remaining = base_model_counts.get(existing_base, 0) - 1 + if remaining > 0: + base_model_counts[existing_base] = remaining + else: + base_model_counts.pop(existing_base, None) + loras[index] = civitai_entry + return + loras.append(civitai_entry) + + def merge_or_append_local(local_entry): + local_id = get_version_id(local_entry) + local_hash = (local_entry.get('hash') or '').lower() + for existing in loras: + existing_id = get_version_id(existing) + existing_hash = (existing.get('hash') or '').lower() + if not ( + (local_id and existing_id == local_id) + or (local_hash and existing_hash == local_hash) + ): + continue + + existing['weight'] = local_entry['weight'] + existing['hash'] = local_entry['hash'] + existing['file_name'] = local_entry['file_name'] + existing['existsLocally'] = True + existing['localPath'] = local_entry['localPath'] + existing['size'] = local_entry['size'] + existing['isDeleted'] = False + if not existing.get('modelId') and local_entry.get('modelId'): + existing['modelId'] = local_entry['modelId'] + if not existing.get('baseModel') and local_entry.get('baseModel'): + existing['baseModel'] = local_entry['baseModel'] + base_model_counts[local_entry['baseModel']] = base_model_counts.get(local_entry['baseModel'], 0) + 1 + thumbnail_url = local_entry.get('thumbnailUrl') + if thumbnail_url and not thumbnail_url.endswith('/images/no-preview.png'): + existing['thumbnailUrl'] = thumbnail_url + return + + if local_entry.get('baseModel'): + base_model = local_entry['baseModel'] + base_model_counts[base_model] = base_model_counts.get(base_model, 0) + 1 + loras.append(local_entry) + + resolved_prompt_basenames = set() + queried_local_basenames = set() + for lora_type, lora_name, lora_weight in prompt_loras.values(): + weight = round(float(lora_weight), 2) + basename_key = (lora_type, normalize_lora_name(lora_name, True)) + matching_resources = [ + lora + for lora in loras[:resource_lora_count] + if lora.get('file_name') + and normalize_lora_name(lora['file_name'], True) == basename_key[1] + and ( + (lora_type == 'hypernet' and str(lora.get('type', '')).casefold() in ('hypernet', 'hypernetwork')) + or (lora_type == 'lora' and str(lora.get('type', '')).casefold() not in ('hypernet', 'hypernetwork')) + ) + ] + if len(prompt_by_basename[basename_key]) == 1 and len(matching_resources) == 1: + matching_resources[0]['weight'] = weight + if basename_key not in hash_basenames: + resolved_prompt_basenames.add(basename_key) + continue + + if basename_key in hash_basenames: + continue + + if not recipe_scanner or lora_type != 'lora': + continue + queried_local_basenames.add(basename_key) + local_lora = await recipe_scanner.get_local_lora(lora_name, recipe_base_model) + if not local_lora: + continue + + local_entry = self.populate_lora_from_local( + make_lora_entry(lora_type, lora_name, weight), + local_lora, + ) + merge_or_append_local(local_entry) + resolved_prompt_basenames.add(basename_key) + + for hash_key, lora_hash in metadata.get("hashes", {}).items(): + if not hash_key.startswith(("lora:", "hypernet:")): + continue + lora_type, lora_name = hash_key.split(':', 1) + basename_key = (lora_type, normalize_lora_name(lora_name, True)) + if basename_key in resolved_prompt_basenames: + continue + + prompt_entries = prompt_by_basename.get(basename_key, []) + weight = prompt_entries[0][1] if len(prompt_entries) == 1 else 1.0 + lora_entry = make_lora_entry(lora_type, lora_name, weight, lora_hash) + + if lora_hash and recipe_scanner and lora_type == 'lora': + local_lora = await recipe_scanner.get_local_lora_by_hash(lora_hash) + if local_lora: + local_entry = self.populate_lora_from_local(lora_entry, local_lora) + merge_or_append_local(local_entry) + continue + + hash_resolved = False + if lora_hash and metadata_provider: + try: + civitai_info = await metadata_provider.get_model_by_hash(lora_hash) + populated_entry = await self.populate_lora_from_civitai( + lora_entry, + civitai_info, + recipe_scanner, + base_model_counts, + lora_hash, + ) + if populated_entry is None: continue - - # Skip entries without a hash value — they can't be - # resolved via CivitAI and would only produce a - # useless "Deleted" entry in the recipe. - if not lora_hash: - continue - - lora_type, lora_name = hash_key.split(':', 1) - - # Get weight from extranet tags if available, else default to 1.0 - weight = lora_weights.get(hash_key, 1.0) - - # Initialize lora entry - lora_entry = { - 'name': lora_name, - 'type': lora_type, # 'lora' or 'hypernet' - 'weight': weight, - 'hash': lora_hash, - 'existsLocally': False, - 'localPath': None, - 'file_name': lora_name, - 'thumbnailUrl': '/loras_static/images/no-preview.png', - 'baseModel': '', - 'size': 0, - 'downloadUrl': '', - 'isDeleted': False - } - - # Try to get info from Civitai - if metadata_provider: - try: - civitai_info = await metadata_provider.get_model_by_hash(lora_hash) - - populated_entry = await self.populate_lora_from_civitai( - lora_entry, - civitai_info, - recipe_scanner, - base_model_counts, - lora_hash - ) - if populated_entry is None: - continue # Skip invalid LoRA types - lora_entry = populated_entry - except Exception as e: - logger.error(f"Error fetching Civitai info for LoRA {lora_name}: {e}") - - loras.append(lora_entry) + lora_entry = populated_entry + hash_resolved = not lora_entry.get('isDeleted') + except Exception as e: + logger.error(f"Error fetching Civitai info for LoRA {lora_name}: {e}") + + if hash_resolved: + merge_or_append_civitai(lora_entry, preserve_existing_weight=not prompt_entries) + continue + + if recipe_scanner and lora_type == 'lora' and basename_key not in queried_local_basenames: + local_lora = await recipe_scanner.get_local_lora(lora_name, recipe_base_model) + if local_lora: + local_entry = self.populate_lora_from_local(lora_entry, local_lora) + merge_or_append_local(local_entry) + continue + + if lora_hash and not resource_lora_count: + loras.append(lora_entry) # Try to get base model from resources or make educated guess base_model = None diff --git a/py/recipes/parsers/comfy.py b/py/recipes/parsers/comfy.py index e1d7251e..678ad075 100644 --- a/py/recipes/parsers/comfy.py +++ b/py/recipes/parsers/comfy.py @@ -31,79 +31,15 @@ class ComfyMetadataParser(RecipeMetadataParser): metadata_provider = await get_default_metadata_provider() data = json.loads(user_comment) - loras = [] - - # Find all LoraLoader nodes - lora_nodes = {k: v for k, v in data.items() if isinstance(v, dict) and v.get('class_type') == 'LoraLoader'} - - # Process each LoraLoader node - for node_id, node in lora_nodes.items(): - if 'inputs' not in node or 'lora_name' not in node['inputs']: - continue - - lora_name = node['inputs'].get('lora_name', '') - - # Parse the URN to extract model ID and version ID - # Format: "urn:air:sdxl:lora:civitai:1107767@1253442" - lora_id_match = re.search(r'civitai:(\d+)@(\d+)', lora_name) - if not lora_id_match: - continue - - model_id = lora_id_match.group(1) - model_version_id = lora_id_match.group(2) - - # Get strength from node inputs - weight = node['inputs'].get('strength_model', 1.0) - - # Initialize lora entry with default values - lora_entry = { - 'id': model_version_id, - 'modelId': model_id, - 'name': f"Lora {model_id}", # Default name - 'version': '', - 'type': 'lora', - 'weight': weight, - 'existsLocally': False, - 'localPath': None, - 'file_name': '', - 'hash': '', - 'thumbnailUrl': '/loras_static/images/no-preview.png', - 'baseModel': '', - 'size': 0, - 'downloadUrl': '', - 'isDeleted': False - } - - # Get additional info from Civitai if metadata provider is available - if metadata_provider: - try: - civitai_info_tuple = await metadata_provider.get_model_version_info(model_version_id) - # Populate lora entry with Civitai info - populated_entry = await self.populate_lora_from_civitai( - lora_entry, - civitai_info_tuple, - recipe_scanner - ) - if populated_entry is None: - continue # Skip invalid LoRA types - lora_entry = populated_entry - except Exception as e: - logger.error(f"Error fetching Civitai info for LoRA: {e}") - - loras.append(lora_entry) - - # Find checkpoint info + checkpoint_nodes = {k: v for k, v in data.items() if isinstance(v, dict) and v.get('class_type') == 'CheckpointLoaderSimple'} checkpoint = None checkpoint_id = None checkpoint_version_id = None - if checkpoint_nodes: - # Get the first checkpoint node checkpoint_node = next(iter(checkpoint_nodes.values())) if 'inputs' in checkpoint_node and 'ckpt_name' in checkpoint_node['inputs']: checkpoint_name = checkpoint_node['inputs']['ckpt_name'] - # Parse checkpoint URN checkpoint_match = re.search(r'civitai:(\d+)@(\d+)', checkpoint_name) if checkpoint_match: checkpoint_id = checkpoint_match.group(1) @@ -115,16 +51,107 @@ class ComfyMetadataParser(RecipeMetadataParser): 'version': '', 'type': 'checkpoint' } - - # Get additional checkpoint info from Civitai if metadata_provider: try: civitai_info_tuple = await metadata_provider.get_model_version_info(checkpoint_version_id) civitai_info, _ = civitai_info_tuple if isinstance(civitai_info_tuple, tuple) else (civitai_info_tuple, None) - # Populate checkpoint with Civitai info checkpoint = await self.populate_checkpoint_from_civitai(checkpoint, civitai_info) except Exception as e: logger.error(f"Error fetching Civitai info for checkpoint: {e}") + + recipe_base_model = checkpoint.get('baseModel') if checkpoint else None + loras = [] + lora_candidates = [] + for node in data.values(): + if not isinstance(node, dict): + continue + + inputs = node.get('inputs') + if not isinstance(inputs, dict): + continue + + if node.get('class_type') == 'LoraLoader': + lora_name = inputs.get('lora_name', '') + if isinstance(lora_name, str) and lora_name: + lora_candidates.append((lora_name, inputs.get('strength_model', 1.0))) + continue + + if node.get('class_type') != 'LoraLoaderLM': + continue + + loras_data = inputs.get('loras', []) + if isinstance(loras_data, dict): + loras_data = loras_data.get('__value__', []) + if isinstance(loras_data, list) and len(loras_data) == 1 and isinstance(loras_data[0], list): + loras_data = loras_data[0] + if not isinstance(loras_data, list): + continue + + for lora in loras_data: + if not isinstance(lora, dict) or not lora.get('active', False) or lora.get('_isDummy', False): + continue + lora_name = lora.get('name', '') + if isinstance(lora_name, str) and lora_name: + lora_candidates.append((lora_name, lora.get('strength', 1.0))) + + for lora_name, weight in lora_candidates: + if isinstance(weight, str): + try: + weight = float(weight) + except ValueError: + weight = 1.0 + lora_id_match = re.search(r'civitai:(\d+)@(\d+)', lora_name) + if lora_id_match: + model_id = lora_id_match.group(1) + model_version_id = lora_id_match.group(2) + entry_name = f"Lora {model_id}" + else: + model_id = 0 + model_version_id = 0 + entry_name = re.split(r'[\\/]', lora_name)[-1] + entry_name = re.sub(r'\.[^.]+$', '', entry_name) + + lora_entry = { + 'id': model_version_id, + 'modelId': model_id, + 'name': entry_name, + 'version': '', + 'type': 'lora', + 'weight': weight, + 'existsLocally': False, + 'localPath': None, + 'file_name': entry_name, + 'hash': '', + 'thumbnailUrl': '/loras_static/images/no-preview.png', + 'baseModel': '', + 'size': 0, + 'downloadUrl': '', + 'isDeleted': False + } + + if lora_id_match: + if metadata_provider: + try: + civitai_info_tuple = await metadata_provider.get_model_version_info(model_version_id) + populated_entry = await self.populate_lora_from_civitai( + lora_entry, + civitai_info_tuple, + recipe_scanner + ) + if populated_entry is None: + continue + lora_entry = populated_entry + except Exception as e: + logger.error(f"Error fetching Civitai info for LoRA: {e}") + else: + if not recipe_scanner: + continue + local_lora = await recipe_scanner.get_local_lora(lora_name, recipe_base_model) + if not local_lora: + continue + lora_entry = self.populate_lora_from_local(lora_entry, local_lora) + + loras.append(lora_entry) # Extract generation parameters gen_params = {} diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index b52f4b28..94d26b08 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -2926,13 +2926,70 @@ class RecipeScanner: return normalized - async def get_local_lora(self, name: str) -> Optional[Dict[str, Any]]: - """Lookup a local LoRA model by name.""" + async def get_local_lora( + self, name: str, base_model: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Lookup an unambiguous local LoRA by name and optional base model.""" if not self._lora_scanner or not name: return None - return await self._lora_scanner.get_model_info_by_name(name) + normalized_name = str(name).replace("\\", "/").casefold() + for extension in (".safetensors", ".ckpt", ".pt", ".bin"): + if normalized_name.endswith(extension): + normalized_name = normalized_name[: -len(extension)] + break + has_path = "/" in normalized_name + basename = normalized_name.rsplit("/", 1)[-1] + + cached_data = await self._lora_scanner.get_cached_data() + matches = [] + for model in cached_data.raw_data: + file_name = str(model.get("file_name") or "").replace("\\", "/") + folder = str(model.get("folder") or "").replace("\\", "/").strip("/") + model_path = f"{folder}/{file_name}" if folder else file_name + for extension in (".safetensors", ".ckpt", ".pt", ".bin"): + if model_path.casefold().endswith(extension): + model_path = model_path[: -len(extension)] + break + if (has_path and model_path.casefold() == normalized_name) or ( + not has_path and model_path.rsplit("/", 1)[-1].casefold() == basename + ): + matches.append(model) + + if len(matches) != 1: + return None + + match = matches[0] + expected_base = str(base_model or "").strip().casefold() + actual_base = str(match.get("base_model") or "").strip().casefold() + if ( + expected_base + and expected_base != "unknown" + and actual_base + and actual_base != "unknown" + and expected_base != actual_base + ): + return None + return match + + async def get_local_lora_by_hash(self, hash_value: str) -> Optional[Dict[str, Any]]: + """Lookup a local LoRA through the scanner's hash index.""" + + if not self._lora_scanner or not hash_value: + return None + + file_path = self._lora_scanner.get_path_by_hash(hash_value) + if not file_path: + return None + + target_path = os.path.normcase(os.path.abspath(file_path)) + cached_data = await self._lora_scanner.get_cached_data() + for model in cached_data.raw_data: + model_path = model.get("file_path") + if model_path and os.path.normcase(os.path.abspath(model_path)) == target_path: + return model + return None async def get_local_checkpoint(self, name: str) -> Optional[Dict[str, Any]]: """Lookup a local checkpoint model by name.""" diff --git a/tests/services/test_automatic_metadata_parser.py b/tests/services/test_automatic_metadata_parser.py index 68ed0511..aed235e4 100644 --- a/tests/services/test_automatic_metadata_parser.py +++ b/tests/services/test_automatic_metadata_parser.py @@ -3,6 +3,40 @@ import pytest from py.recipes.parsers.automatic import AutomaticMetadataParser +class LocalRecipeScanner: + class LoraScanner: + @staticmethod + def has_hash(model_hash): + return False + + def __init__(self, models): + self.models = models + self.queries = [] + self.hash_queries = [] + self._lora_scanner = self.LoraScanner() + + async def get_local_lora(self, name, base_model=None): + self.queries.append(name) + return self.models.get(name) + + async def get_local_lora_by_hash(self, hash_value): + self.hash_queries.append(hash_value) + return next((model for model in self.models.values() if model.get("sha256") == hash_value), None) + + +def local_lora(file_name="local_only"): + return { + "file_path": f"/models/loras/styles/{file_name}.safetensors", + "file_name": file_name, + "model_name": "Local Only", + "sha256": "a" * 64, + "size": 123456, + "base_model": "Flux.1 D", + "preview_url": f"/models/loras/styles/{file_name}.preview.png", + "civitai": None, + } + + @pytest.mark.asyncio async def test_parse_metadata_extracts_checkpoint_from_civitai_resources(monkeypatch): checkpoint_info = { @@ -132,6 +166,218 @@ async def test_parse_metadata_merges_lora_hashes_over_empty_hashes_json(monkeypa assert "UnusedLora" not in lora_names, "UnusedLora should have been skipped" +@pytest.mark.asyncio +async def test_parse_metadata_resolves_local_lora_with_empty_hash(monkeypatch): + async def fake_metadata_provider(): + class Provider: + async def get_model_by_hash(self, model_hash): + raise AssertionError("Local and empty-hash LoRAs must not query Civitai") + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local_only": local_lora()}) + metadata_text = ( + "portrait \n" + "Steps: 20, Sampler: Euler, CFG scale: 7, Seed: 1, " + 'Hashes: {"lora:local_only": "", "lora:missing": ""}' + ) + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert len(result["loras"]) == 1 + entry = result["loras"][0] + assert entry["name"] == "Local Only" + assert entry["file_name"] == "local_only" + assert entry["weight"] == 0.65 + assert entry["hash"] == "a" * 64 + assert entry["localPath"].endswith("local_only.safetensors") + assert entry["size"] == 123456 + assert entry["baseModel"] == "Flux.1 D" + assert entry["existsLocally"] is True + assert entry["isDeleted"] is False + assert scanner.queries == ["local_only", "missing"] + + +@pytest.mark.asyncio +async def test_parse_metadata_resolves_prompt_lora_without_hashes(monkeypatch): + async def fake_metadata_provider(): + return None + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + model = local_lora() + scanner = LocalRecipeScanner({"styles/local_only": model}) + metadata_text = "portrait \nSteps: 20, Seed: 1" + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert len(result["loras"]) == 1 + assert result["loras"][0]["weight"] == 0.7 + assert result["loras"][0]["hash"] == "a" * 64 + assert scanner.queries == ["styles/local_only"] + + +@pytest.mark.asyncio +async def test_parse_metadata_prefers_hash_over_colliding_local_name(monkeypatch): + remote_info = { + "id": 100, + "modelId": 200, + "model": {"name": "Hash Match", "type": "LORA"}, + "name": "v1", + "files": [{"type": "Model", "primary": True, "name": "hash_match.safetensors", "hashes": {"SHA256": "b" * 64}}], + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_by_hash(self, model_hash): + assert model_hash == "deadbeef00" + return remote_info, None + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local_only": local_lora()}) + metadata_text = ( + "portrait \n" + "Steps: 20, Seed: 1, " + 'Hashes: {"lora:local_only": "deadbeef00"}' + ) + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert len(result["loras"]) == 1 + assert result["loras"][0]["id"] == 100 + assert result["loras"][0]["weight"] == 0.8 + assert scanner.queries == [] + assert scanner.hash_queries == ["deadbeef00"] + + +@pytest.mark.asyncio +async def test_parse_metadata_falls_back_to_name_when_hash_is_unresolved(monkeypatch): + async def fake_metadata_provider(): + class Provider: + async def get_model_by_hash(self, model_hash): + return None, "Model not found" + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local_only": local_lora()}) + metadata_text = ( + "portrait \nSteps: 20, Seed: 1, " + 'Hashes: {"lora:local_only": "deadbeef00"}' + ) + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert result["loras"][0]["hash"] == "a" * 64 + assert scanner.queries == ["local_only"] + + +@pytest.mark.asyncio +async def test_parse_metadata_uses_prompt_weight_for_civitai_resource(monkeypatch): + remote_info = { + "id": 100, + "modelId": 200, + "model": {"name": "local_only", "type": "LORA"}, + "name": "v1", + "files": [ + { + "type": "Model", + "primary": True, + "name": "remote_file.safetensors", + "hashes": {"SHA256": "b" * 64}, + } + ], + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + assert version_id == 100 + return remote_info, None + + async def get_model_by_hash(self, model_hash): + raise AssertionError("The Civitai resource should not be fetched again by hash") + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local_only": local_lora()}) + metadata_text = ( + "portrait \n" + "Steps: 20, Seed: 1, " + 'Civitai resources: [{"type":"lora","modelVersionId":100,"modelName":"local_only"}]' + ) + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert len(result["loras"]) == 2 + assert [entry["file_name"] for entry in result["loras"]] == ["remote_file", "local_only"] + assert [entry["weight"] for entry in result["loras"]] == [0.35, 0.6] + assert result["loras"][1]["existsLocally"] is True + assert scanner.queries == ["local_only"] + + +@pytest.mark.asyncio +async def test_parse_metadata_keeps_mixed_local_and_civitai_loras(monkeypatch): + remote_info = { + "id": 100, + "modelId": 200, + "model": {"name": "Remote LoRA", "type": "LORA"}, + "name": "v1", + "files": [ + { + "type": "Model", + "primary": True, + "name": "remote.safetensors", + "hashes": {"SHA256": "b" * 64}, + } + ], + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_by_hash(self, model_hash): + assert model_hash == "bbbbbbbbbb" + return remote_info, None + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.automatic.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local_only": local_lora()}) + metadata_text = ( + "portrait \n" + "Steps: 20, Seed: 1, " + 'Hashes: {"lora:local_only": "", "lora:remote": "bbbbbbbbbb"}' + ) + + result = await AutomaticMetadataParser().parse_metadata(metadata_text, scanner) + + assert [entry["name"] for entry in result["loras"]] == ["Local Only", "Remote LoRA"] + assert [entry["weight"] for entry in result["loras"]] == [0.6, 0.9] + assert result["loras"][0]["existsLocally"] is True + assert result["loras"][1]["existsLocally"] is False + + @pytest.mark.asyncio async def test_parse_metadata_extracts_checkpoint_from_model_hash(monkeypatch): checkpoint_info = { diff --git a/tests/services/test_comfy_metadata_parser.py b/tests/services/test_comfy_metadata_parser.py index dac489ab..25136f56 100644 --- a/tests/services/test_comfy_metadata_parser.py +++ b/tests/services/test_comfy_metadata_parser.py @@ -2,6 +2,38 @@ import pytest import json from py.recipes.parsers.comfy import ComfyMetadataParser + +class LocalRecipeScanner: + class LoraScanner: + @staticmethod + def has_hash(model_hash): + return False + + def __init__(self, models): + self.models = models + self.queries = [] + self.base_models = [] + self._lora_scanner = self.LoraScanner() + + async def get_local_lora(self, name, base_model=None): + self.queries.append(name) + self.base_models.append(base_model) + return self.models.get(name) + + +def local_lora(file_name): + return { + "file_path": f"/models/loras/{file_name}.safetensors", + "file_name": file_name.rsplit("/", 1)[-1], + "model_name": file_name.rsplit("/", 1)[-1], + "sha256": file_name[0] * 64, + "size": 4096, + "base_model": "SDXL 1.0", + "preview_url": "", + "civitai": None, + } + + @pytest.mark.asyncio async def test_parse_metadata_without_loras(monkeypatch): checkpoint_info = { @@ -84,6 +116,140 @@ async def test_parse_metadata_without_loras(monkeypatch): assert result["gen_params"]["size"] == "1024x1024" assert result["from_comfy_metadata"] is True +@pytest.mark.asyncio +async def test_parse_metadata_resolves_standard_and_manager_local_loras(monkeypatch): + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + raise AssertionError("Local LoRAs must not query Civitai") + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.comfy.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({ + "styles/standard.safetensors": local_lora("standard"), + "manager": local_lora("manager"), + }) + metadata_json = { + "1": { + "class_type": "LoraLoader", + "inputs": { + "lora_name": "styles/standard.safetensors", + "strength_model": 0.55, + }, + }, + "2": { + "class_type": "LoraLoaderLM", + "inputs": { + "loras": { + "__value__": [ + {"name": "manager", "strength": "0.80", "active": True}, + {"name": "disabled", "strength": 1.0, "active": False}, + {"name": "dummy", "strength": 1.0, "active": True, "_isDummy": True}, + ] + } + }, + }, + } + + result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner) + + assert [entry["file_name"] for entry in result["loras"]] == ["standard", "manager"] + assert [entry["weight"] for entry in result["loras"]] == [0.55, 0.8] + assert all(isinstance(entry["weight"], float) for entry in result["loras"]) + assert all(entry["existsLocally"] is True for entry in result["loras"]) + assert all(entry["isDeleted"] is False for entry in result["loras"]) + assert scanner.queries == ["styles/standard.safetensors", "manager"] + + +@pytest.mark.asyncio +async def test_parse_metadata_defaults_malformed_weight_and_passes_checkpoint_base_model(monkeypatch): + checkpoint_info = { + "id": 456, + "modelId": 123, + "model": {"name": "Checkpoint", "type": "checkpoint"}, + "name": "v1", + "baseModel": "SDXL 1.0", + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + return checkpoint_info, None + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.comfy.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"style": local_lora("style")}) + metadata_json = { + "1": {"class_type": "LoraLoader", "inputs": {"lora_name": "style", "strength_model": "invalid"}}, + "2": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "civitai:123@456"}}, + } + + result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner) + + assert result["loras"][0]["weight"] == 1.0 + assert scanner.base_models == ["SDXL 1.0"] + + +@pytest.mark.asyncio +async def test_parse_metadata_keeps_civitai_urn_with_local_lora(monkeypatch): + remote_info = { + "id": 456, + "modelId": 123, + "model": {"name": "Remote LoRA", "type": "LORA"}, + "name": "v1", + "files": [ + { + "type": "Model", + "primary": True, + "name": "remote.safetensors", + "hashes": {"SHA256": "c" * 64}, + } + ], + } + + async def fake_metadata_provider(): + class Provider: + async def get_model_version_info(self, version_id): + assert version_id == "456" + return remote_info, None + + return Provider() + + monkeypatch.setattr( + "py.recipes.parsers.comfy.get_default_metadata_provider", + fake_metadata_provider, + ) + scanner = LocalRecipeScanner({"local": local_lora("local")}) + metadata_json = { + "1": { + "class_type": "LoraLoader", + "inputs": {"lora_name": "local", "strength_model": 0.4}, + }, + "2": { + "class_type": "LoraLoader", + "inputs": { + "lora_name": "urn:air:sdxl:lora:civitai:123@456", + "strength_model": 0.9, + }, + }, + } + + result = await ComfyMetadataParser().parse_metadata(json.dumps(metadata_json), scanner) + + assert [entry["name"] for entry in result["loras"]] == ["local", "Remote LoRA"] + assert [entry["weight"] for entry in result["loras"]] == [0.4, 0.9] + assert result["loras"][0]["existsLocally"] is True + assert result["loras"][1]["id"] == 456 + + @pytest.mark.asyncio async def test_parse_metadata_without_extra_metadata(monkeypatch): async def fake_metadata_provider(): diff --git a/tests/services/test_recipe_scanner.py b/tests/services/test_recipe_scanner.py index 0c3c75da..015920be 100644 --- a/tests/services/test_recipe_scanner.py +++ b/tests/services/test_recipe_scanner.py @@ -107,6 +107,35 @@ def recipe_scanner(tmp_path: Path, monkeypatch): settings_manager_module.reset_settings_manager() +@pytest.mark.asyncio +async def test_local_lora_lookup_requires_unambiguous_name_and_matching_base_model(recipe_scanner): + scanner, stub = recipe_scanner + models = [ + { + "file_name": "style.safetensors", + "folder": "sd15", + "file_path": "/models/loras/sd15/style.safetensors", + "sha256": "a" * 64, + "base_model": "SD 1.5", + }, + { + "file_name": "style.safetensors", + "folder": "sdxl", + "file_path": "/models/loras/sdxl/style.safetensors", + "sha256": "b" * 64, + "base_model": "SDXL 1.0", + }, + ] + stub._cache.raw_data = models + stub._hash_meta["b" * 64] = {"path": models[1]["file_path"]} + + assert await scanner.get_local_lora("style") is None + assert await scanner.get_local_lora("sdxl/style.safetensors", "SDXL 1.0") is models[1] + assert await scanner.get_local_lora("sdxl/style.safetensors", "SD 1.5") is None + assert await scanner.get_local_lora("other/style.safetensors") is None + assert await scanner.get_local_lora_by_hash("b" * 64) is models[1] + + def test_recipes_dir_uses_custom_settings_path(tmp_path: Path, monkeypatch): RecipeScanner._instance = None settings_manager_module.reset_settings_manager()