From 9ff0752ef16ea19ead6601904c6c64b8437fcb96 Mon Sep 17 00:00:00 2001 From: JJ Date: Tue, 28 Oct 2025 12:58:20 +0000 Subject: [PATCH] lora tests l --- __init__.py | 3 + nodes/image_nodes.py | 130 ++++---------- nodes/lora_change.py | 299 +++++++++++++++++++++++++++++++++ nodes/lora_combine.py | 383 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 714 insertions(+), 101 deletions(-) create mode 100644 nodes/lora_change.py create mode 100644 nodes/lora_combine.py diff --git a/__init__.py b/__init__.py index 1368450..9a726c3 100644 --- a/__init__.py +++ b/__init__.py @@ -10,6 +10,9 @@ from .nodes.decode_overlap import * from .nodes.power_spline_editor import * from .nodes.lora_extract_v2 import * +from .nodes.lora_change import * +from .nodes.lora_combine import * + from .nodes.draw_shapes import * from .nodes.draw_joints import * from .nodes.draw_image import * diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py index ff7504d..6a57cf9 100644 --- a/nodes/image_nodes.py +++ b/nodes/image_nodes.py @@ -449,8 +449,8 @@ class WanScaleAB: return { "required": { "image_A": ("IMAGE",), - "A1size": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 16}), - "A2scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}), + "A2size": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 16}), + "A1scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}), "scaling_method": (["area","lanczos","bilinear", "bicubic" ,"nearest" ], {"default": "area"}), "match_to": (["A_crop", "A_stretch", "B_crop", "B_stretch",], {"default": "A_crop"}), }, @@ -464,46 +464,23 @@ class WanScaleAB: FUNCTION = "scale_images" CATEGORY = "WanVideoWrapper_QQ/image" - def scale_images(self, image_A, A1size, A2scale, scaling_method, match_to, image_B=None): + def scale_images(self, image_A, A2size, A1scale, scaling_method, match_to, image_B=None): # If B_stretch is selected and image_B is present, use modular approach if image_B is not None and match_to == "B_stretch": # Step 1: Calculate image_B2 and image_B1 dimensions using image_A's methods - b_s1, b_s2 = self._calculate_s1_s2(image_B, A1size) + b_s1, b_s2 = self._calculate_s1_s2(image_B, A2size) # Get image_B dimensions for aspect ratio _, b_h, b_w, _ = image_B.shape - # Calculate image_B1 dimensions with aspect ratio preservation - # Special case: when A2scale=1.0, B1 should match B2 dimensions exactly - if abs(A2scale - 1.0) < 1e-6: - # Use exactly the same dimensions as B2 (b_s1, b_s2) + # Calculate image_B1 dimensions as scaled version of B2 + # Special case: when A1scale=1.0, B1 should match B2 dimensions exactly + if abs(A1scale - 1.0) < 1e-6: b_target_h = b_s1 b_target_w = b_s2 else: - # Calculate the base target dimensions - if b_h >= b_w: - base_height = b_s1 * A2scale - base_width = b_s2 * A2scale - else: - base_height = b_s2 * A2scale - base_width = b_s1 * A2scale - - # Calculate aspect ratio of the original image_B - original_aspect = b_w / b_h - - # Calculate target dimensions that preserve aspect ratio - if original_aspect > (base_width / base_height): - # Image is wider than target aspect, match width - b_target_w = base_width - b_target_h = base_width / original_aspect - else: - # Image is taller than or equal to target aspect, match height - b_target_h = base_height - b_target_w = base_height * original_aspect - - # Round up to multiples of 16 - b_target_h = self._round_up_to_multiple(b_target_h, 16) - b_target_w = self._round_up_to_multiple(b_target_w, 16) + b_target_h = self._round_up_to_multiple(b_s1 * A1scale, 16) + b_target_w = self._round_up_to_multiple(b_s2 * A1scale, 16) # Step 2: Now reverse the process - use these dimensions for image_A # Calculate image_A2 and image_A1 using image_B's calculated dimensions @@ -516,43 +493,18 @@ class WanScaleAB: elif image_B is not None and match_to == "B_crop": # B_crop should work like B_stretch for dimension calculations but use cropping # Step 1: Calculate image_B2 and image_B1 dimensions using image_B's methods (like B_stretch) - b_s1, b_s2 = self._calculate_s1_s2(image_B, A1size) + b_s1, b_s2 = self._calculate_s1_s2(image_B, A2size) # Step 2: Scale and crop image_A to match image_B2's dimensions (reversed from A_crop) image_A2 = self._scale_and_crop_to_match(image_A, b_s1, b_s2, scaling_method) - # Step 3: Calculate image_B1 dimensions with aspect ratio preservation - _, b_h, b_w, _ = image_B.shape - # Special case: when A2scale=1.0, B1 should match B2 dimensions exactly - if abs(A2scale - 1.0) < 1e-6: - # Use exactly the same dimensions as B2 (b_s1, b_s2) + # Step 3: Calculate image_B1 dimensions as scaled version of B2 + if abs(A1scale - 1.0) < 1e-6: b_target_h = b_s1 b_target_w = b_s2 else: - # Calculate the base target dimensions - if b_h >= b_w: - base_height = b_s1 * A2scale - base_width = b_s2 * A2scale - else: - base_height = b_s2 * A2scale - base_width = b_s1 * A2scale - - # Calculate aspect ratio of the original image_B - original_aspect = b_w / b_h - - # Calculate target dimensions that preserve aspect ratio - if original_aspect > (base_width / base_height): - # Image is wider than target aspect, match width - b_target_w = base_width - b_target_h = base_width / original_aspect - else: - # Image is taller than or equal to target aspect, match height - b_target_h = base_height - b_target_w = base_height * original_aspect - - # Round up to multiples of 16 - b_target_h = self._round_up_to_multiple(b_target_h, 16) - b_target_w = self._round_up_to_multiple(b_target_w, 16) + b_target_h = self._round_up_to_multiple(b_s1 * A1scale, 16) + b_target_w = self._round_up_to_multiple(b_s2 * A1scale, 16) # Step 4: Scale image_B to preserve aspect ratio while fitting within target dimensions image_B1 = self._scale_and_crop_to_match(image_B, int(b_target_h), int(b_target_w), scaling_method) @@ -564,43 +516,19 @@ class WanScaleAB: else: # Original logic for other match_to options (A_crop, A_stretch, etc.) # Calculate s1 and s2 for image_A2 - s1, s2 = self._calculate_s1_s2(image_A, A1size) + s1, s2 = self._calculate_s1_s2(image_A, A2size) # Get original dimensions of image_A _, h, w, _ = image_A.shape - # Process image_A1 with aspect ratio preservation - # Special case: when A2scale=1.0, A1 should match A2 dimensions exactly - if abs(A2scale - 1.0) < 1e-6: - # Use exactly the same dimensions as A2 (s1, s2) + # Process image_A1 as a direct scaled version of A2 (s1, s2) + # Special case: when A1scale=1.0, A1 should match A2 dimensions exactly + if abs(A1scale - 1.0) < 1e-6: target_height = s1 target_width = s2 else: - # Calculate the base target dimensions - if h >= w: - base_height = s1 * A2scale - base_width = s2 * A2scale - else: - base_height = s2 * A2scale - base_width = s1 * A2scale - - # Calculate aspect ratio of the original image - original_aspect = w / h - - # Calculate target dimensions that preserve aspect ratio - # We'll fit the image within the base dimensions while maintaining aspect ratio - if original_aspect > (base_width / base_height): - # Image is wider than target aspect, match width - target_width = base_width - target_height = base_width / original_aspect - else: - # Image is taller than or equal to target aspect, match height - target_height = base_height - target_width = base_height * original_aspect - - # Find closest target dimensions divisible by 16 (rounding up) - target_height = self._round_up_to_multiple(target_height, 16) - target_width = self._round_up_to_multiple(target_width, 16) + target_height = self._round_up_to_multiple(s1 * A1scale, 16) + target_width = self._round_up_to_multiple(s2 * A1scale, 16) # Scale image_A1 and image_A2 image_A1 = self._scale_image(image_A, int(target_height), int(target_width), scaling_method) @@ -623,8 +551,8 @@ class WanScaleAB: # Get original dimensions _, h, w, _ = image_B.shape # Calculate new dimensions based on scale - new_h = int(h * A2scale) - new_w = int(w * A2scale) + new_h = int(h * A1scale) + new_w = int(w * A1scale) image_B1 = self._scale_image(image_B, new_h, new_w, scaling_method) image_B2 = self._scale_image(image_B, new_h, new_w, scaling_method) else: @@ -634,11 +562,11 @@ class WanScaleAB: return (image_A1, image_A2, image_B1, image_B2) - def _calculate_s1_s2(self, image_A, A1size): + def _calculate_s1_s2(self, image_A, A2size): """ Calculate s1 and s2 based on the special scaling rules. - s1: Find the closest number to A1size that can divide the largest side of image_A fully once, + s1: Find the closest number to A2size that can divide the largest side of image_A fully once, then adjust to be divisible by 16 (rounding up if needed) s2: Proportionally scale the smallest side to match s1, then adjust to be divisible by 16 (rounding up) """ @@ -654,13 +582,13 @@ class WanScaleAB: smallest_side = h # Step 1: Find s1 - # Find how many times A1size fits into the largest side - times_fit = largest_side // A1size + # Find how many times A2size fits into the largest side + times_fit = largest_side // A2size if times_fit == 0: times_fit = 1 # Ensure at least once - # Calculate the closest number that could divide the largest side by A1size fully once - base_s1 = A1size * times_fit + # Calculate the closest number that could divide the largest side by A2size fully once + base_s1 = A2size * times_fit # Find the closest number to base_s1 that is divisible by 16 s1 = self._round_up_to_multiple(base_s1, 16) @@ -780,4 +708,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageRadialZoomBlur_GPU": "Image Radial Zoom Blur (GPU)", "ImageBlend_GPU": "Image Blend (GPU)", "WanScaleAB": "Wan Scale AB" -} \ No newline at end of file +} diff --git a/nodes/lora_change.py b/nodes/lora_change.py new file mode 100644 index 0000000..e12608a --- /dev/null +++ b/nodes/lora_change.py @@ -0,0 +1,299 @@ +import os +import torch +import comfy.utils +import folder_paths + + +class LoraReduceRank: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}), + "new_rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1, "tooltip": "The new rank to resize the LoRA. Acts as max rank when using dynamic_method."}), + "dynamic_method": (["disabled", "sv_ratio", "sv_cumulative", "sv_fro"], {"default": "disabled", "tooltip": "Method to use for dynamically determining new alphas and dims"}), + "dynamic_param": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Method to use for dynamically determining new alphas and dims"}), + "output_dtype": (["match_original", "fp16", "bf16", "fp32"], {"default": "match_original", "tooltip": "Data type to save the LoRA as."}), + "verbose": ("BOOLEAN", {"default": True}), + }, + + } + RETURN_TYPES = () + FUNCTION = "save" + OUTPUT_NODE = True + EXPERIMENTAL = True + DESCRIPTION = "Resize a LoRA model by reducing it's rank. Based on kohya's sd-scripts." + + CATEGORY = "KJNodes/lora" + + def save(self, lora_name, new_rank, output_dtype, dynamic_method, dynamic_param, verbose): + + lora_path = folder_paths.get_full_path("loras", lora_name) + lora_sd, metadata = comfy.utils.load_torch_file(lora_path, return_metadata=True) + + if output_dtype == "fp16": + save_dtype = torch.float16 + elif output_dtype == "bf16": + save_dtype = torch.bfloat16 + elif output_dtype == "fp32": + save_dtype = torch.float32 + elif output_dtype == "match_original": + first_weight_key = next(k for k in lora_sd if k.endswith(".weight") and isinstance(lora_sd[k], torch.Tensor)) + save_dtype = lora_sd[first_weight_key].dtype + + # Note: resize_lora_model is defined in lora_nodes.py; this class copy is for organization only. + # Import lazily to avoid circular imports if needed. + from .lora_nodes import resize_lora_model, device # type: ignore + + new_lora_sd = {} + for k, v in lora_sd.items(): + new_lora_sd[k.replace(".default", "")] = v + del lora_sd + print("Resizing Lora...") + output_sd, old_dim, new_alpha, rank_list = resize_lora_model(new_lora_sd, new_rank, save_dtype, device, dynamic_method, dynamic_param, verbose) + + # update metadata + if metadata is None: + metadata = {} + + comment = metadata.get("ss_training_comment", "") + + if dynamic_method == "disabled": + metadata["ss_training_comment"] = f"dimension is resized from {old_dim} to {new_rank}; {comment}" + metadata["ss_network_dim"] = str(new_rank) + metadata["ss_network_alpha"] = str(new_alpha) + else: + metadata["ss_training_comment"] = f"Dynamic resize with {dynamic_method}: {dynamic_param} from {old_dim}; {comment}" + metadata["ss_network_dim"] = "Dynamic" + metadata["ss_network_alpha"] = "Dynamic" + + # cast to save_dtype before calculating hashes + for key in list(output_sd.keys()): + value = output_sd[key] + if type(value) == torch.Tensor and value.dtype.is_floating_point and value.dtype != save_dtype: + output_sd[key] = value.to(save_dtype) + + output_filename_prefix = "loras/" + lora_name + + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(output_filename_prefix, self.output_dir) + output_dtype_str = f"_{output_dtype}" if output_dtype != "match_original" else "" + import numpy as np + average_rank = str(int(np.mean(rank_list))) + rank_str = new_rank if dynamic_method == "disabled" else f"dynamic_{average_rank}" + output_checkpoint = f"{filename.replace('.safetensors', '')}_resized_from_{old_dim}_to_{rank_str}{output_dtype_str}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + print(f"Saving resized LoRA to {output_checkpoint}") + + comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=metadata) + return {} + + +class LoraRemapStrength: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "LoRA to remap."}), + "reference_strength": ("FLOAT", {"default": 0.15, "min": 1e-6, "max": 10.0, "step": 0.01, "tooltip": "Existing effective strength you typically use (e.g., 0.15)."}), + "target_strength": ("FLOAT", {"default": 1.0, "min": 1e-6, "max": 10.0, "step": 0.01, "tooltip": "Desired equivalent strength (e.g., 1.0)."}), + "create_missing_alpha": ("BOOLEAN", {"default": True, "tooltip": "Create per-module .alpha if missing (uses rank as base)."}), + "output_dtype": (["match_original", "fp16", "bf16", "fp32"], {"default": "match_original", "tooltip": "Optional cast before save."}), + } + } + + RETURN_TYPES = () + FUNCTION = "save" + OUTPUT_NODE = True + CATEGORY = "KJNodes/lora" + + def save(self, lora_name, reference_strength, target_strength, create_missing_alpha, output_dtype): + print(f"[LoraRemapStrength DEBUG] Starting with inputs:") + print(f" - lora_name: {lora_name}") + print(f" - reference_strength: {reference_strength}") + print(f" - target_strength: {target_strength}") + print(f" - create_missing_alpha: {create_missing_alpha}") + print(f" - output_dtype: {output_dtype}") + + lora_path = folder_paths.get_full_path("loras", lora_name) + print(f"[LoraRemapStrength DEBUG] Loading LoRA from: {lora_path}") + + try: + lora_sd, metadata = comfy.utils.load_torch_file(lora_path, return_metadata=True) + print(f"[LoraRemapStrength DEBUG] Successfully loaded LoRA with {len(lora_sd)} keys") + if metadata: + print(f"[LoraRemapStrength DEBUG] Metadata keys: {list(metadata.keys())}") + else: + print(f"[LoraRemapStrength DEBUG] No metadata found") + except Exception as e: + print(f"[LoraRemapStrength ERROR] Failed to load LoRA: {e}") + return {} + + # compute scale factor to make target behave like reference + factor = float(target_strength) / float(reference_strength) + print(f"[LoraRemapStrength DEBUG] Computed scaling factor: {factor}") + + # Determine save dtype + if output_dtype == "match_original": + # try to find a floating weight to infer dtype + first_weight_key = next((k for k in lora_sd if k.endswith(".weight") and isinstance(lora_sd[k], torch.Tensor)), None) + save_dtype = lora_sd[first_weight_key].dtype if first_weight_key is not None else torch.float16 + print(f"[LoraRemapStrength DEBUG] Match original dtype: {save_dtype} (from key: {first_weight_key})") + elif output_dtype == "fp16": + save_dtype = torch.float16 + print(f"[LoraRemapStrength DEBUG] Using fp16 dtype") + elif output_dtype == "bf16": + save_dtype = torch.bfloat16 + print(f"[LoraRemapStrength DEBUG] Using bf16 dtype") + else: + save_dtype = torch.float32 + print(f"[LoraRemapStrength DEBUG] Using fp32 dtype") + + # Helper to get module prefix and weight name variant + def split_lora_key(key): + if key.endswith(".weight"): + parts = key.split(".") + if len(parts) >= 3: + mod = parts[-2] + if mod in ("lora_down", "lora_up", "lora_A", "lora_B", "down", "up"): + return ".".join(parts[:-2]), mod, True + else: + parts = key.split(".") + if len(parts) >= 2 and parts[-1] in ("lora_down", "lora_up", "lora_A", "lora_B", "down", "up"): + return ".".join(parts[:-1]), parts[-1], False + return None, None, None + + # Collect modules and detect ranks + modules = {} + processed_keys = [] + for k, v in lora_sd.items(): + pref, mod, has_weight = split_lora_key(k) + if pref is None: + continue + processed_keys.append(k) + m = modules.setdefault(pref, {"down": None, "up": None, "alpha": None}) + # track alpha + alpha_key = pref + ".alpha" + if alpha_key in lora_sd: + m["alpha"] = alpha_key + # track weights + if (mod in ("lora_down", "lora_A", "down")): + m["down"] = k + elif (mod in ("lora_up", "lora_B", "up")): + m["up"] = k + + print(f"[LoraRemapStrength DEBUG] Found {len(modules)} LoRA modules") + print(f"[LoraRemapStrength DEBUG] Processed {len(processed_keys)} LoRA keys") + + # Scale existing alphas; create missing if requested + alpha_scaled_count = 0 + alpha_created_count = 0 + for pref, info in modules.items(): + alpha_key = info["alpha"] + if alpha_key is not None and isinstance(lora_sd[alpha_key], torch.Tensor): + try: + original_alpha = lora_sd[alpha_key].clone() + lora_sd[alpha_key] = (lora_sd[alpha_key].to(torch.float32) * factor).to(save_dtype) + alpha_scaled_count += 1 + print(f"[LoraRemapStrength DEBUG] Scaled alpha for {pref}: {original_alpha} -> {lora_sd[alpha_key]}") + except Exception as e: + print(f"[LoraRemapStrength DEBUG] Error scaling alpha for {pref}: {e}") + lora_sd[alpha_key] = torch.tensor(float(lora_sd[alpha_key]) * factor, dtype=save_dtype) + alpha_scaled_count += 1 + elif create_missing_alpha and info["down"] is not None: + w = lora_sd[info["down"]] + rank = w.shape[0] if isinstance(w, torch.Tensor) and w.ndim >= 2 else 1 + lora_sd[pref + ".alpha"] = torch.tensor(float(rank) * factor, dtype=save_dtype) + alpha_created_count += 1 + print(f"[LoraRemapStrength DEBUG] Created alpha for {pref}: rank={rank}, alpha={lora_sd[pref + '.alpha']}") + + print(f"[LoraRemapStrength DEBUG] Scaled {alpha_scaled_count} existing alphas") + print(f"[LoraRemapStrength DEBUG] Created {alpha_created_count} missing alphas") + + # Cast weights to save_dtype for consistency + for key in list(lora_sd.keys()): + val = lora_sd[key] + if isinstance(val, torch.Tensor) and val.dtype.is_floating_point and val.dtype != save_dtype: + lora_sd[key] = val.to(save_dtype) + + # Update metadata + if metadata is None: + metadata = {} + comment = metadata.get("ss_training_comment", "") + metadata["ss_training_comment"] = f"alpha scaled by x{factor:.4f} to remap {reference_strength} -> {target_strength}; {comment}" + + # Build output filename - save next to original LoRA + base_name = os.path.basename(lora_name).replace(".safetensors", "") + suffix = f"_remap_{reference_strength}_to_{target_strength}" + + # Get the directory of the original LoRA file + original_lora_dir = os.path.dirname(lora_path) + + print(f"[LoraRemapStrength DEBUG] Building output path:") + print(f" - base_name: {base_name}") + print(f" - suffix: {suffix}") + print(f" - original_lora_dir: {original_lora_dir}") + print(f" - original_lora_path: {lora_path}") + + try: + # Create output filename in the same directory as the original LoRA + # Add a counter to avoid overwriting existing files + counter = 1 + while True: + output_filename = f"{base_name}{suffix}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(original_lora_dir, output_filename) + if not os.path.exists(output_checkpoint): + break + counter += 1 + + print(f"[LoraRemapStrength DEBUG] Path generation results:") + print(f" - output_checkpoint: {output_checkpoint}") + print(f" - counter: {counter}") + + # Verify the directory exists (it should since it contains the original LoRA) + os.makedirs(original_lora_dir, exist_ok=True) + print(f"[LoraRemapStrength DEBUG] Confirmed output directory exists: {original_lora_dir}") + + except Exception as e: + print(f"[LoraRemapStrength ERROR] Failed to generate output path: {e}") + return {} + + # Update metadata + if metadata is None: + metadata = {} + comment = metadata.get("ss_training_comment", "") + metadata["ss_training_comment"] = f"alpha scaled by x{factor:.4f} to remap {reference_strength} -> {target_strength}; {comment}" + print(f"[LoraRemapStrength DEBUG] Updated metadata comment") + + try: + print(f"[LoraRemapStrength DEBUG] Attempting to save LoRA file...") + comfy.utils.save_torch_file(lora_sd, output_checkpoint, metadata=metadata) + print(f"[LoraRemapStrength SUCCESS] LoRA saved successfully to: {output_checkpoint}") + + # Verify the file was actually created + if os.path.exists(output_checkpoint): + file_size = os.path.getsize(output_checkpoint) + print(f"[LoraRemapStrength DEBUG] File verification: EXISTS, size: {file_size} bytes") + else: + print(f"[LoraRemapStrength ERROR] File verification: FILE NOT FOUND at {output_checkpoint}") + + except Exception as e: + print(f"[LoraRemapStrength ERROR] Failed to save LoRA file: {e}") + return {} + + return {} + +NODE_CLASS_MAPPINGS = { + "LoraReduceRank": LoraReduceRank, + "LoraRemapStrength": LoraRemapStrength +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoraReduceRank": "Lora Reduce Rank", + "LoraRemapStrength": "Lora Remap Strength" +} \ No newline at end of file diff --git a/nodes/lora_combine.py b/nodes/lora_combine.py new file mode 100644 index 0000000..4c0c34a --- /dev/null +++ b/nodes/lora_combine.py @@ -0,0 +1,383 @@ +import os +import torch +import comfy.utils +import folder_paths +import logging +import numpy as np +from typing import Dict, List, Tuple, Optional + +device = comfy.model_management.get_torch_device() + +def combine_lora_weights(lora_sds: List[Dict], blend_weights: List[float], target_rank: int, + save_dtype: torch.dtype, verbose: bool = False) -> Tuple[Dict, List[int]]: + """ + Combine multiple LoRA state dictionaries with rank extension using SVD. + + Args: + lora_sds: List of LoRA state dictionaries + blend_weights: List of blend weights for each LoRA + target_rank: Target rank for the combined LoRA + save_dtype: Data type for the combined weights + verbose: Whether to print verbose information + + Returns: + Combined LoRA state dictionary and list of actual ranks used + """ + if len(lora_sds) != len(blend_weights): + raise ValueError("Number of LoRAs must match number of blend weights") + + if len(lora_sds) == 0: + raise ValueError("At least one LoRA must be provided") + + # Normalize blend weights + total_weight = sum(blend_weights) + if total_weight == 0: + raise ValueError("Sum of blend weights cannot be zero") + blend_weights = [w / total_weight for w in blend_weights] + + combined_sd = {} + rank_list = [] + + # Get all unique layer prefixes from all LoRAs + all_layer_prefixes = set() + for lora_sd in lora_sds: + for key in lora_sd.keys(): + if key.endswith((".lora_up.weight", ".lora_down.weight", ".lora_A.weight", ".lora_B.weight")): + prefix = key.rsplit('.', 2)[0] + all_layer_prefixes.add(prefix) + + if verbose: + print(f"Found {len(all_layer_prefixes)} unique layers to process") + + # Process each layer + for layer_prefix in sorted(all_layer_prefixes): + try: + # Collect up and down weights from all LoRAs for this layer + up_weights = [] + down_weights = [] + alphas = [] + + for lora_idx, lora_sd in enumerate(lora_sds): + up_key = f"{layer_prefix}.lora_up.weight" + down_key = f"{layer_prefix}.lora_down.weight" + alpha_key = f"{layer_prefix}.alpha" + + # Try alternative LoRA naming conventions + if up_key not in lora_sd: + up_key = f"{layer_prefix}.lora_B.weight" + if down_key not in lora_sd: + down_key = f"{layer_prefix}.lora_A.weight" + + if up_key in lora_sd and down_key in lora_sd: + up_weights.append(lora_sd[up_key].to(device) * blend_weights[lora_idx]) + down_weights.append(lora_sd[down_key].to(device) * blend_weights[lora_idx]) + + if alpha_key in lora_sd: + alphas.append(lora_sd[alpha_key].item() if isinstance(lora_sd[alpha_key], torch.Tensor) else lora_sd[alpha_key]) + + if not up_weights: + if verbose: + print(f"Skipping layer {layer_prefix}: no compatible weights found") + continue + + # Combine weights additively + combined_up = torch.stack(up_weights).sum(dim=0) + combined_down = torch.stack(down_weights).sum(dim=0) + + # Apply SVD-based rank extension + extended_up, extended_down = extend_rank_with_svd( + combined_up, combined_down, target_rank, verbose + ) + + # Store combined weights + combined_sd[f"{layer_prefix}.lora_up.weight"] = extended_up.cpu().contiguous().to(save_dtype) + combined_sd[f"{layer_prefix}.lora_down.weight"] = extended_down.cpu().contiguous().to(save_dtype) + + # Handle alpha values + if alphas: + # Blend alphas proportionally + combined_alpha = sum(alpha * weight for alpha, weight in zip(alphas, blend_weights[:len(alphas)])) + combined_sd[f"{layer_prefix}.alpha"] = torch.tensor(combined_alpha, dtype=save_dtype) + + rank_list.append(target_rank) + + if verbose: + print(f"Combined layer {layer_prefix}: original shapes {combined_up.shape, combined_down.shape} -> extended rank {target_rank}") + + except Exception as e: + logging.warning(f"Failed to combine layer {layer_prefix}: {e}") + continue + + return combined_sd, rank_list + +def extend_rank_with_svd(up_weight: torch.Tensor, down_weight: torch.Tensor, + target_rank: int, verbose: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Extend LoRA rank using SVD decomposition. + + Args: + up_weight: Up projection weight matrix + down_weight: Down projection weight matrix + target_rank: Target rank for extension + verbose: Whether to print verbose information + + Returns: + Extended up and down weight matrices + """ + # Combine the LoRA matrices by multiplication + # Original LoRA: W + up @ down + # We'll decompose the combined effect and reconstruct with target rank + combined_effect = up_weight @ down_weight + + # Handle different tensor shapes (conv2d vs linear) + if combined_effect.ndim == 4: + # Conv2D case: flatten spatial dimensions + out_dim, in_dim, kh, kw = combined_effect.shape + combined_flat = combined_effect.view(out_dim, -1) + is_conv = True + spatial_shape = (kh, kw) + else: + # Linear case + out_dim, in_dim = combined_effect.shape + combined_flat = combined_effect + is_conv = False + spatial_shape = None + + # SVD decomposition + try: + U, S, Vh = torch.linalg.svd(combined_flat.float(), full_matrices=False) + except Exception as e: + if verbose: + print(f"SVD failed, using fallback method: {e}") + # Fallback: pad with zeros if SVD fails + return fallback_rank_extension(up_weight, down_weight, target_rank, is_conv, spatial_shape) + + # Ensure we don't exceed target dimensions + actual_rank = min(target_rank, len(S), out_dim, in_dim) + + if actual_rank < target_rank and verbose: + print(f"Warning: Target rank {target_rank} exceeds matrix dimensions, using {actual_rank}") + + # Reconstruct with target rank + U_reduced = U[:, :actual_rank] + S_reduced = S[:actual_rank] + Vh_reduced = Vh[:actual_rank, :] + + # Split the reconstructed matrix back into up and down matrices + # We'll distribute the singular values between up and down + sqrt_S = torch.sqrt(S_reduced) + + if is_conv: + # Conv2D case + new_up = U_reduced @ torch.diag(sqrt_S) + new_down = torch.diag(sqrt_S) @ Vh_reduced + + # Reshape back to conv2d format + new_up = new_up.view(out_dim, actual_rank, 1, 1) + new_down = new_down.view(actual_rank, in_dim, kh, kw) + else: + # Linear case + new_up = U_reduced @ torch.diag(sqrt_S) + new_down = torch.diag(sqrt_S) @ Vh_reduced + + if verbose: + # Calculate preservation ratio + original_norm = torch.norm(combined_flat).item() + reconstructed_flat = new_up @ new_down if not is_conv else (new_up @ new_down).view(out_dim, -1) + reconstructed_norm = torch.norm(reconstructed_flat).item() + preservation = reconstructed_norm / original_norm if original_norm > 0 else 0 + print(f"Rank extension: {combined_flat.shape} -> rank {actual_rank}, preservation: {preservation:.2%}") + + return new_up, new_down + +def fallback_rank_extension(up_weight: torch.Tensor, down_weight: torch.Tensor, + target_rank: int, is_conv: bool, spatial_shape: Optional[Tuple]) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Fallback method for rank extension when SVD fails. + """ + out_dim = up_weight.shape[0] + in_dim = down_weight.shape[-1] if is_conv else down_weight.shape[1] + current_rank = up_weight.shape[1] if is_conv else up_weight.shape[1] + + # Can't reduce rank with fallback, only extend + if target_rank <= current_rank: + return up_weight, down_weight + + if is_conv: + # Conv2D case + kh, kw = spatial_shape + # Pad up weight + padding_up = torch.zeros(out_dim, target_rank - current_rank, 1, 1, + device=up_weight.device, dtype=up_weight.dtype) + new_up = torch.cat([up_weight, padding_up], dim=1) + + # Pad down weight + padding_down = torch.zeros(target_rank - current_rank, in_dim, kh, kw, + device=down_weight.device, dtype=down_weight.dtype) + new_down = torch.cat([down_weight, padding_down], dim=0) + else: + # Linear case + # Pad up weight + padding_up = torch.zeros(out_dim, target_rank - current_rank, + device=up_weight.device, dtype=up_weight.dtype) + new_up = torch.cat([up_weight, padding_up], dim=1) + + # Pad down weight + padding_down = torch.zeros(target_rank - current_rank, in_dim, + device=down_weight.device, dtype=down_weight.dtype) + new_down = torch.cat([down_weight, padding_down], dim=0) + + return new_up, new_down + +class LoraSmartCombine: + """ + Smart LoRA combination node that can merge multiple LoRAs while extending their rank. + """ + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "lora_1_name": (folder_paths.get_filename_list("loras"), {"tooltip": "First LoRA to combine"}), + "lora_2_name": (folder_paths.get_filename_list("loras"), {"tooltip": "Second LoRA to combine"}), + "lora_3_name": (["None"] + folder_paths.get_filename_list("loras"), {"tooltip": "Third LoRA (optional)"}), + "lora_4_name": (["None"] + folder_paths.get_filename_list("loras"), {"tooltip": "Fourth LoRA (optional)"}), + + "target_rank": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1, + "tooltip": "Target rank for the combined LoRA"}), + + "blend_weight_1": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, + "tooltip": "Blend weight for first LoRA"}), + "blend_weight_2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, + "tooltip": "Blend weight for second LoRA"}), + "blend_weight_3": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, + "tooltip": "Blend weight for third LoRA"}), + "blend_weight_4": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, + "tooltip": "Blend weight for fourth LoRA"}), + + "alpha_handling": (["preserve_original", "blend_proportional", "recalculate_from_rank"], + {"default": "blend_proportional", + "tooltip": "How to handle alpha values when combining"}), + + "output_dtype": (["match_original", "fp16", "bf16", "fp32"], + {"default": "match_original", "tooltip": "Data type for output LoRA"}), + + "verbose": ("BOOLEAN", {"default": True, "tooltip": "Print detailed combination information"}), + }, + } + + RETURN_TYPES = () + FUNCTION = "combine" + OUTPUT_NODE = True + CATEGORY = "KJNodes/lora" + DESCRIPTION = "Smartly combine multiple LoRAs with rank extension using SVD decomposition" + + def combine(self, lora_1_name, lora_2_name, lora_3_name, lora_4_name, + target_rank, blend_weight_1, blend_weight_2, blend_weight_3, blend_weight_4, + alpha_handling, output_dtype, verbose): + + # Collect LoRA names and weights + lora_names = [lora_1_name, lora_2_name] + blend_weights = [blend_weight_1, blend_weight_2] + + if lora_3_name != "None": + lora_names.append(lora_3_name) + blend_weights.append(blend_weight_3) + if lora_4_name != "None": + lora_names.append(lora_4_name) + blend_weights.append(blend_weight_4) + + # Validate inputs + if len(lora_names) < 2: + raise ValueError("At least two LoRAs must be selected") + + if verbose: + print(f"\n{'='*60}") + print(f"Smart LoRA Combination Started") + print(f"LoRAs to combine: {len(lora_names)}") + print(f"Target rank: {target_rank}") + print(f"Blend weights: {blend_weights}") + print(f"Alpha handling: {alpha_handling}") + print(f"{'='*60}\n") + + # Load LoRA state dictionaries + lora_sds = [] + original_dtype = None + + for lora_name in lora_names: + lora_path = folder_paths.get_full_path("loras", lora_name) + lora_sd, metadata = comfy.utils.load_torch_file(lora_path, return_metadata=True) + lora_sds.append(lora_sd) + + # Determine original dtype from first LoRA + if original_dtype is None and output_dtype == "match_original": + first_weight_key = next((k for k in lora_sd if k.endswith(".weight") and isinstance(lora_sd[k], torch.Tensor)), None) + if first_weight_key is not None: + original_dtype = lora_sd[first_weight_key].dtype + + # Determine output dtype + if output_dtype == "match_original": + save_dtype = original_dtype or torch.float16 + elif output_dtype == "fp16": + save_dtype = torch.float16 + elif output_dtype == "bf16": + save_dtype = torch.bfloat16 + else: # fp32 + save_dtype = torch.float32 + + if verbose: + print(f"Using output dtype: {save_dtype}") + + # Combine LoRAs + combined_sd, rank_list = combine_lora_weights( + lora_sds, blend_weights, target_rank, save_dtype, verbose + ) + + if not combined_sd: + raise ValueError("No layers were successfully combined") + + # Create metadata + metadata = { + "ss_training_comment": f"Smart combination of {len(lora_names)} LoRAs with rank {target_rank}", + "ss_network_dim": str(target_rank), + "ss_network_alpha": str(target_rank), # Use target rank as alpha + "combination_method": "smart_svd_extension", + "source_loras": lora_names, + "blend_weights": blend_weights, + "alpha_handling": alpha_handling + } + + # Generate output filename + lora_base_names = [os.path.basename(name).replace('.safetensors', '') for name in lora_names] + combination_name = "_plus_".join(lora_base_names[:2]) + if len(lora_base_names) > 2: + combination_name += f"_plus_{len(lora_base_names)-2}_more" + + output_filename_prefix = f"loras/{combination_name}_smart_combined" + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(output_filename_prefix, self.output_dir) + + output_checkpoint = f"{filename}_rank_{target_rank}_{output_dtype}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + + # Save combined LoRA + comfy.utils.save_torch_file(combined_sd, output_checkpoint, metadata=metadata) + + if verbose: + print(f"\n{'='*60}") + print(f"Smart LoRA Combination Completed") + print(f"Output saved to: {output_checkpoint}") + print(f"Combined {len(combined_sd)//2} layers") + print(f"Average rank: {np.mean(rank_list):.1f}") + print(f"{'='*60}\n") + + return {} + +NODE_CLASS_MAPPINGS = { + "LoraSmartCombine": LoraSmartCombine +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoraSmartCombine": "Lora Smart Combine" +} \ No newline at end of file