From ce55b6d7f3ede9700a6cf8e107e607922a47c304 Mon Sep 17 00:00:00 2001 From: Oliver Switzer Date: Mon, 15 Sep 2025 23:20:43 -0400 Subject: [PATCH] feat(prompt-composer): add WAN 2.2 high/low LoRA pairing support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements automatic detection and pairing of WAN 2.2 high/low LoRAs: - Add is_wan_2_2_lora() detection function in metadata utils - Implement robust case-preserving pairing algorithm - Integrate pairing into video LoRA selection workflow - Respect max_video_loras limit (pairs count as single units) - Ensure paired LoRAs have complete data structure - Add comprehensive test coverage (8 new tests) Examples: - Wan22-I2V-HIGH-Robot → Wan22-I2V-LOW-Robot - character_high → character_low - motion_highnoise → motion_lownoise Output: 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- nodes/lora_metadata_utils.py | 42 +++++++++++ nodes/prompt_composer_node.py | 118 +++++++++++++++++++++++++++++- tests/test_lora_metadata_utils.py | 48 ++++++++++++ tests/test_prompt_composer.py | 82 +++++++++++++++++++++ 4 files changed, 289 insertions(+), 1 deletion(-) diff --git a/nodes/lora_metadata_utils.py b/nodes/lora_metadata_utils.py index 8b08c00..fee36c3 100644 --- a/nodes/lora_metadata_utils.py +++ b/nodes/lora_metadata_utils.py @@ -465,6 +465,48 @@ def classify_lora_type(metadata: dict[str, Any]) -> str: return "unknown" +def is_wan_2_2_lora(metadata: dict[str, Any]) -> bool: + """ + Determine if a LoRA is specifically for WAN 2.2 (requires high/low pairing). + + Args: + metadata: LoRA metadata dictionary + + Returns: + True if this is a WAN 2.2 LoRA that needs high/low pairing + """ + if not metadata: + return False + + # Check base_model field for WAN 2.2 indicators + base_model = metadata.get("base_model", "").lower() + wan_2_2_indicators = ["wan2.2", "wan 2.2", "wanv2.2", "wan v2.2", "wan_2.2"] + if any(indicator in base_model for indicator in wan_2_2_indicators): + return True + + # Check civitai base model + civitai = metadata.get("civitai") + if isinstance(civitai, dict): + civitai_base = civitai.get("baseModel", "").lower() + if any(indicator in civitai_base for indicator in wan_2_2_indicators): + return True + + # Check model name and description for WAN 2.2 mentions + model_name = metadata.get("model_name", "").lower() + if any(indicator in model_name for indicator in wan_2_2_indicators): + return True + + # Check civitai model name + if isinstance(civitai, dict): + model = civitai.get("model") + if isinstance(model, dict): + civitai_model_name = model.get("name", "").lower() + if any(indicator in civitai_model_name for indicator in wan_2_2_indicators): + return True + + return False + + def extract_recommended_weight(metadata: dict[str, Any]) -> float: """ Extract recommended weight/strength for a LoRA from its metadata. diff --git a/nodes/prompt_composer_node.py b/nodes/prompt_composer_node.py index db70d48..57f9fbe 100644 --- a/nodes/prompt_composer_node.py +++ b/nodes/prompt_composer_node.py @@ -7,7 +7,7 @@ composes prompts with optimal LoRA combinations, weights, and trigger words. import json import re -from typing import Any +from typing import Any, Optional from .logging_utils import log, log_debug, log_error from .lora_metadata_utils import ( @@ -219,6 +219,113 @@ This node log_error(f"Failed to initialize embeddings: {e}") return False + def _find_wan_lora_pair(self, selected_lora_name: str) -> Optional[str]: + """ + Find the matching HIGH/LOW pair for a WAN LoRA using simple word replacement. + + Args: + selected_lora_name: Name of the selected WAN LoRA + + Returns: + Name of the matching pair LoRA, or None if no pair found + """ + import re + + name_lower = selected_lora_name.lower() + + # Check if it contains "high" or "low" (case insensitive) + has_high = "high" in name_lower + has_low = "low" in name_lower + + if not (has_high or has_low): + return None # Not a high/low variant + + # Simple replacement: swap "high" with "low" and vice versa (preserve case) + if has_high: + # Replace preserving case: HIGH->LOW, High->Low, high->low + def replace_high(match): + original = match.group(0) + if original.isupper(): + return "LOW" + elif original.istitle(): + return "Low" + else: + return "low" + + pair_name = re.sub(r"high", replace_high, selected_lora_name, flags=re.IGNORECASE) + else: # has_low + # Replace preserving case: LOW->HIGH, Low->High, low->high + def replace_low(match): + original = match.group(0) + if original.isupper(): + return "HIGH" + elif original.istitle(): + return "High" + else: + return "high" + + pair_name = re.sub(r"low", replace_low, selected_lora_name, flags=re.IGNORECASE) + + # Check if the pair exists in available LoRAs + if pair_name in self._lora_database: + return pair_name + + return None + + def _apply_wan_2_2_pairing( + self, video_loras: list[dict[str, Any]], max_count: int + ) -> list[dict[str, Any]]: + """ + Post-process video LoRA selection to add WAN 2.2 high/low pairs. + Respects the max_count limit by treating pairs as single units. + + Args: + video_loras: List of selected video LoRAs + max_count: Maximum number of LoRA "units" (pairs count as 1 unit) + + Returns: + Expanded list with high/low pairs added for WAN 2.2 LoRAs + """ + from .lora_metadata_utils import is_wan_2_2_lora + + expanded_loras = [] + units_used = 0 + + for lora in video_loras: + if units_used >= max_count: + log(f"⚠️ Reached max_video_loras limit ({max_count}), skipping remaining LoRAs") + break + + expanded_loras.append(lora) # Always add the original + + # Check if this is a WAN 2.2 LoRA + metadata = lora.get("metadata") + if metadata and is_wan_2_2_lora(metadata): + # Try to find the matching pair + pair_name = self._find_wan_lora_pair(lora["name"]) + + if pair_name and pair_name in self._lora_database: + # Add the pair LoRA with same relevance score and proper structure + pair_db_info = self._lora_database[pair_name] + pair_metadata = pair_db_info["metadata"] + pair_info = { + "name": pair_name, + "metadata": pair_metadata, + "relevance_score": lora.get("relevance_score", 0.0), + "recommended_weight": extract_recommended_weight(pair_metadata), + "trigger_words": pair_db_info.get("trigger_words", []), + "type": "video", + } + expanded_loras.append(pair_info) + + log(f"🔗 WAN 2.2 pairing: {lora['name']} + {pair_name}") + else: + log(f"⚠️ WAN 2.2 LoRA {lora['name']} has no matching pair") + + units_used += 1 # Each original LoRA (with or without pair) counts as 1 unit + + return expanded_loras + def _find_relevant_loras( self, scene_description: str, @@ -775,6 +882,15 @@ This node content_boost, wan_lora_dir_path, ) + + # Apply WAN 2.2 high/low pairing + original_count = len(video_loras) + video_loras = self._apply_wan_2_2_pairing(video_loras, max_video_loras) + if len(video_loras) > original_count: + log( + f"WAN 2.2 pairing expanded video LoRAs from {original_count} to {len(video_loras)}" + ) + video_names = [lora.get("name", "unknown") for lora in video_loras] log(f"Found {len(video_loras)} video LoRAs: {video_names}") diff --git a/tests/test_lora_metadata_utils.py b/tests/test_lora_metadata_utils.py index 613b9e8..670ee64 100644 --- a/tests/test_lora_metadata_utils.py +++ b/tests/test_lora_metadata_utils.py @@ -17,6 +17,7 @@ from nodes.lora_metadata_utils import ( extract_recommended_weight, get_lora_trigger_words, is_video_lora, + is_wan_2_2_lora, load_lora_metadata, ) @@ -301,6 +302,53 @@ class TestNewMetadataFunctions(unittest.TestCase): metadata = {"modelDescription": "Use weight 0.05"} self.assertEqual(extract_recommended_weight(metadata), 0.8) # Should use default + def test_is_wan_2_2_lora_base_model(self): + """is_wan_2_2_lora should detect WAN 2.2 from base_model field.""" + # Test various WAN 2.2 indicators in base_model + metadata = {"base_model": "WAN2.2"} + self.assertTrue(is_wan_2_2_lora(metadata)) + + metadata = {"base_model": "wan 2.2 video model"} + self.assertTrue(is_wan_2_2_lora(metadata)) + + metadata = {"base_model": "WAN v2.2"} + self.assertTrue(is_wan_2_2_lora(metadata)) + + # Test non-WAN 2.2 models + metadata = {"base_model": "WAN2.1"} + self.assertFalse(is_wan_2_2_lora(metadata)) + + metadata = {"base_model": "SDXL"} + self.assertFalse(is_wan_2_2_lora(metadata)) + + def test_is_wan_2_2_lora_civitai_model(self): + """is_wan_2_2_lora should detect WAN 2.2 from civitai fields.""" + # Test civitai baseModel + metadata = {"civitai": {"baseModel": "WAN2.2"}} + self.assertTrue(is_wan_2_2_lora(metadata)) + + # Test civitai model name + metadata = {"civitai": {"model": {"name": "Character WAN2.2 LoRA"}}} + self.assertTrue(is_wan_2_2_lora(metadata)) + + # Test non-WAN 2.2 + metadata = {"civitai": {"baseModel": "WAN"}} + self.assertFalse(is_wan_2_2_lora(metadata)) + + def test_is_wan_2_2_lora_model_name(self): + """is_wan_2_2_lora should detect WAN 2.2 from model_name field.""" + metadata = {"model_name": "Dancing Girl WAN2.2"} + self.assertTrue(is_wan_2_2_lora(metadata)) + + metadata = {"model_name": "Regular LoRA"} + self.assertFalse(is_wan_2_2_lora(metadata)) + + def test_is_wan_2_2_lora_empty_metadata(self): + """is_wan_2_2_lora should handle empty or None metadata.""" + self.assertFalse(is_wan_2_2_lora(None)) + self.assertFalse(is_wan_2_2_lora({})) + self.assertFalse(is_wan_2_2_lora({"unrelated_field": "value"})) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_prompt_composer.py b/tests/test_prompt_composer.py index 2fc01e8..dd22196 100644 --- a/tests/test_prompt_composer.py +++ b/tests/test_prompt_composer.py @@ -221,6 +221,88 @@ class TestPromptComposerNode(unittest.TestCase): self.assertTrue(result) self.assertTrue(self.node._embeddings_initialized) + def test_find_wan_lora_pair_high_to_low(self): + """_find_wan_lora_pair should find matching low pair for high LoRA.""" + # Setup mock database with pairs + self.node._lora_database = { + "character_high": {"metadata": {}}, + "character_low": {"metadata": {}}, + "motion_highnoise": {"metadata": {}}, + "motion_lownoise": {"metadata": {}}, + } + + # Test various high->low patterns + self.assertEqual(self.node._find_wan_lora_pair("character_high"), "character_low") + self.assertEqual(self.node._find_wan_lora_pair("motion_highnoise"), "motion_lownoise") + + def test_find_wan_lora_pair_dash_patterns(self): + """_find_wan_lora_pair should handle dash-separated HIGH/LOW patterns.""" + # Setup mock database with dash patterns + self.node._lora_database = { + "Wan22-I2V-HIGH-Dancing_Robot": {"metadata": {}}, + "Wan22-I2V-LOW-Dancing_Robot": {"metadata": {}}, + } + + # Test dash-separated HIGH/LOW pattern + result = self.node._find_wan_lora_pair("Wan22-I2V-HIGH-Dancing_Robot") + self.assertEqual(result, "Wan22-I2V-LOW-Dancing_Robot") + + def test_find_wan_lora_pair_low_to_high(self): + """_find_wan_lora_pair should find matching high pair for low LoRA.""" + # Setup mock database with pairs + self.node._lora_database = { + "anime-girl-high": {"metadata": {}}, + "anime-girl-low": {"metadata": {}}, + "dancer_high": {"metadata": {}}, + "dancer_low": {"metadata": {}}, + } + + # Test various low->high patterns + self.assertEqual(self.node._find_wan_lora_pair("anime-girl-low"), "anime-girl-high") + self.assertEqual(self.node._find_wan_lora_pair("dancer_low"), "dancer_high") + + def test_find_wan_lora_pair_no_pair(self): + """_find_wan_lora_pair should return None when no pair exists.""" + # Setup mock database with single LoRA + self.node._lora_database = { + "character_high": {"metadata": {}}, + "standalone_lora": {"metadata": {}}, + } + + # Test missing pair + self.assertIsNone(self.node._find_wan_lora_pair("character_high")) # No "character_low" + self.assertIsNone(self.node._find_wan_lora_pair("standalone_lora")) # Not high/low variant + + def test_apply_wan_2_2_pairing(self): + """_apply_wan_2_2_pairing should add pairs for WAN 2.2 LoRAs.""" + # Setup mock database and video LoRAs + self.node._lora_database = { + "wan22_dancer_high": { + "metadata": {"base_model": "WAN2.2"}, + "trigger_words": ["dance"], + }, + "wan22_dancer_low": { + "metadata": {"base_model": "WAN2.2"}, + "trigger_words": ["dance"], + }, + } + + video_loras = [ + { + "name": "wan22_dancer_high", + "metadata": {"base_model": "WAN2.2"}, + "relevance_score": 0.85, + } + ] + + result = self.node._apply_wan_2_2_pairing(video_loras, max_count=3) + + # Should have original + pair + self.assertEqual(len(result), 2) + self.assertEqual(result[0]["name"], "wan22_dancer_high") + self.assertEqual(result[1]["name"], "wan22_dancer_low") + self.assertEqual(result[1]["relevance_score"], 0.85) # Same score as original + if __name__ == "__main__": unittest.main()