From 7f97af84673feec580cfc96f17eb11a0d51653a2 Mon Sep 17 00:00:00 2001 From: JJ Date: Wed, 25 Mar 2026 04:46:01 +0000 Subject: [PATCH] AB overlap --- __init__.py | 1 + nodes/lora_merge_ltx.py | 591 ++++++++++++++++ nodes/video_nodes_extra.py | 1186 ++++++++++++++++++++++++++++++--- utility/ltx/__init__.py | 48 ++ utility/ltx/core/__init__.py | 23 + utility/ltx/core/merger.py | 436 ++++++++++++ utility/ltx/core/ops.py | 552 +++++++++++++++ utility/ltx/utils/__init__.py | 22 + utility/ltx/utils/io.py | 65 ++ utility/ltx/utils/naming.py | 203 ++++++ utility/ltx/utils/progress.py | 81 +++ 11 files changed, 3124 insertions(+), 84 deletions(-) create mode 100644 nodes/lora_merge_ltx.py create mode 100644 utility/ltx/__init__.py create mode 100644 utility/ltx/core/__init__.py create mode 100644 utility/ltx/core/merger.py create mode 100644 utility/ltx/core/ops.py create mode 100644 utility/ltx/utils/__init__.py create mode 100644 utility/ltx/utils/io.py create mode 100644 utility/ltx/utils/naming.py create mode 100644 utility/ltx/utils/progress.py diff --git a/__init__.py b/__init__.py index 55f080e..71013a5 100644 --- a/__init__.py +++ b/__init__.py @@ -13,6 +13,7 @@ from .nodes.video_inpaint import * from .nodes.matanyone2 import * from .nodes.lora_extract_v2 import * +from .nodes.lora_merge_ltx import * from .nodes.lora_change import * from .nodes.lora_combine import * from .nodes.math_nodes import * diff --git a/nodes/lora_merge_ltx.py b/nodes/lora_merge_ltx.py new file mode 100644 index 0000000..7804116 --- /dev/null +++ b/nodes/lora_merge_ltx.py @@ -0,0 +1,591 @@ +""" +LTX2 LoRA Select - Simplified LTX2 LoRA loader with metadata pass-through. +Designed for daisy-chaining multiple LoRA loaders with combined metadata tracking. +""" + +import os +import sys +from pathlib import Path + +# Add parent directory to sys.path to import utility module +parent_dir = Path(__file__).resolve().parent.parent +if str(parent_dir) not in sys.path: + sys.path.insert(0, str(parent_dir)) + +import torch +import json +from safetensors.torch import save_file + +import comfy.lora +import comfy.utils +import comfy.sd +import folder_paths + +from utility.ltx import merge_lora_state_dicts, count_layer_types + + +class LoadModelMeta: + """ + Loads a checkpoint model and stores the full directory path in metadata. + Similar to checkpoint loader but with metadata tracking. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "checkpoint_name": ( + folder_paths.get_filename_list("checkpoints"), + {"tooltip": "The name of the checkpoint file."}, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "load_checkpoint" + CATEGORY = "KJNodes/ltxv" + + def load_checkpoint(self, checkpoint_name): + """ + Load checkpoint and store the full path in model_options. + """ + # Get full path to checkpoint + checkpoint_path = folder_paths.get_full_path("checkpoints", checkpoint_name) + + # Load the checkpoint + model = comfy.sd.load_checkpoint_guess_config( + checkpoint_path, output_vae=True, output_clip=True + )[0] + + # Store checkpoint path in model_options for LTX2LoRASelect to read + if not hasattr(model, "model_options"): + model.model_options = {} + + # Store the full path as checkpoint_name + model.model_options["checkpoint_name"] = checkpoint_path + + return (model,) + + +class LTX2LoRASelect: + """ + Simplified LTX2 LoRA loader with per-block strength control. + Outputs model with attached metadata for daisy-chaining multiple LoRA loaders. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "lora_name": ( + folder_paths.get_filename_list("loras"), + {"tooltip": "The name of the LoRA file."}, + ), + "model": ( + "MODEL", + {"tooltip": "The diffusion model to apply LoRA to."}, + ), + "strength_model": ( + "FLOAT", + { + "default": 1.0, + "min": -100.0, + "max": 100.0, + "step": 0.01, + "tooltip": "Overall LoRA strength multiplier.", + }, + ), + "video": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength for video attention layers.", + }, + ), + "video_to_audio": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength for video to audio cross-attention layers.", + }, + ), + "audio": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength for audio attention layers.", + }, + ), + "audio_to_video": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength for audio to video cross-attention layers.", + }, + ), + "other": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Strength for layers not caught by other layer filters.", + }, + ), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply_lora" + CATEGORY = "KJNodes/ltxv" + + def apply_lora( + self, + model, + lora_name, + strength_model, + video, + video_to_audio, + audio, + audio_to_video, + other, + ): + """ + Load and apply LoRA with layer-specific strength control. + Attaches metadata to the model for daisy-chaining multiple LoRA loaders. + """ + # Load LoRA file + lora_path = folder_paths.get_full_path("loras", lora_name) + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + + # Get rank info for metadata + rank = "unknown" + weight_key = next((key for key in lora.keys() if key.endswith("weight")), None) + if weight_key: + rank = str(lora[weight_key].shape[0]) + + # Build key map for current model + key_map = {} + if model is not None: + key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) + + # Load LoRA weights + loaded = comfy.lora.load_lora(lora, key_map) + + keys_to_delete = [] + + # Apply layer-based attention strength filtering + for key in list(loaded.keys()): + if key in keys_to_delete: + continue + + key_str = ( + key + if isinstance(key, str) + else (key[0] if isinstance(key, tuple) else str(key)) + ) + + # Determine strength multiplier based on layer name + strength_multiplier = None + + # Video to audio cross-attention (check first - most specific) + if "video_to_audio_attn" in key_str: + strength_multiplier = video_to_audio + # Audio to video cross-attention + elif "audio_to_video_attn" in key_str: + strength_multiplier = audio_to_video + # Audio layers + elif "audio_attn" in key_str or "audio_ff.net" in key_str: + strength_multiplier = audio + # Video layers (check last - most general) + elif "attn" in key_str or "ff.net" in key_str: + strength_multiplier = video + # Everything else not caught by above filters + else: + strength_multiplier = other + + # Apply strength or mark for deletion + if strength_multiplier is not None: + if strength_multiplier == 0: + keys_to_delete.append(key) + elif strength_multiplier != 1.0: + value = loaded[key] + if hasattr(value, "weights"): + weights_list = list(value.weights) + # Handle case where alpha (weights[2]) might be None + current_alpha = ( + weights_list[2] if weights_list[2] is not None else 1.0 + ) + weights_list[2] = current_alpha * strength_multiplier + loaded[key].weights = tuple(weights_list) + + for key in keys_to_delete: + if key in loaded: + del loaded[key] + + # Apply patches to model + if model is not None: + new_modelpatcher = model.clone() + new_modelpatcher.add_patches(loaded, strength_model) + + # Get connected model info for metadata + model_id = "unknown" + checkpoint_path = "" + + if hasattr(model, "model_options"): + # Try to get checkpoint path from model_options (set by LoadModelMeta) + checkpoint_path = model.model_options.get("checkpoint_name", "") + + # Get model type as fallback + if hasattr(model.model, "model_config"): + model_id = str(type(model.model).__name__) + elif hasattr(model.model, "load_device"): + model_id = f"{type(model.model).__name__}" + + elif hasattr(model, "model"): + if hasattr(model.model, "model_config"): + model_id = str(type(model.model).__name__) + elif hasattr(model.model, "load_device"): + model_id = f"{type(model.model).__name__}" + + # Create and attach metadata for daisy-chaining + lora_metadata = { + "lora_name": lora_name, + "strength_model": strength_model, + "video": video, + "video_to_audio": video_to_audio, + "audio": audio, + "audio_to_video": audio_to_video, + "other": other, + "rank": rank, + "applied_keys_count": len(loaded), + "connected_model": model_id, + "checkpoint_path": checkpoint_path, + } + + # Attach metadata to model_options for daisy-chaining + if not hasattr(new_modelpatcher, "model_options"): + new_modelpatcher.model_options = {} + + # Preserve checkpoint_name from input model + if checkpoint_path: + new_modelpatcher.model_options["checkpoint_name"] = checkpoint_path + + if "ltx2_lora_chain" not in new_modelpatcher.model_options: + new_modelpatcher.model_options["ltx2_lora_chain"] = [] + + # Append this LoRA's metadata to the chain + new_modelpatcher.model_options["ltx2_lora_chain"].append(lora_metadata) + else: + new_modelpatcher = None + + return (new_modelpatcher,) + + +class LoraMergeLTX: + """ + Merges multiple LoRAs that were applied via LTX2LoRASelect into a single LoRA file. + + This node reads the ltx2_lora_chain metadata from a model that has had + multiple LoRAs applied, then merges them into a single safetensors file. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_with_loras": ( + "MODEL", + {"tooltip": "The model with LoRAs applied (from LTX2LoRASelect)."}, + ), + "output_prefix": ( + "STRING", + { + "default": "merged_lora", + "tooltip": "Prefix for the merged LoRA filename.", + }, + ), + "rank": ( + "INT", + { + "default": 64, + "min": 1, + "max": 512, + "step": 8, + "tooltip": "Target rank of the merged LoRA.", + }, + ), + "output_dir": ( + "STRING", + { + "default": "loras", + "tooltip": "Output directory (relative to ComfyUI root). Use 'loras' for auto-discovery.", + }, + ), + "merge_on_model": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Merge on model weights (more accurate) vs pure LoRA SVD merge (faster).", + }, + ), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("merged_lora_path",) + FUNCTION = "merge_loras" + CATEGORY = "KJNodes/ltxv" + OUTPUT_NODE = True + + def merge_loras(self, model_with_loras, output_prefix, rank, output_dir, merge_on_model): + """ + Merge multiple LoRAs from ltx2_lora_chain into a single safetensors file. + + The merge process: + 1. Extract the base model's state dict + 2. For each LoRA in the chain, compute its effective weight contribution + (considering layer-specific strengths) + 3. Sum all contributions to get merged delta weights + 4. Apply SVD decomposition to factorize into up/down matrices at target rank + 5. Save as safetensors with metadata + """ + # Validate input + if model_with_loras is None: + raise ValueError("No model provided for merging") + + if not hasattr(model_with_loras, "model_options"): + raise ValueError("Model has no model_options attribute") + + lora_chain = model_with_loras.model_options.get("ltx2_lora_chain", []) + if not lora_chain: + raise ValueError( + "No LoRA chain found in model options. Apply LoRAs first using LTX2LoRASelect." + ) + + print(f"\n{'=' * 60}") + print(f"LTX2 LoRA Merge - Starting") + print(f"{'=' * 60}") + print(f"Found {len(lora_chain)} LoRA(s) to merge:") + for i, lora_info in enumerate(lora_chain, 1): + lora_name = lora_info.get('lora_name', 'unknown') + strength_model = lora_info.get('strength_model', 0) + layer_strs = [ + f"v:{lora_info.get('video', 1.0):.2f}", + f"a:{lora_info.get('audio', 1.0):.2f}", + f"v>a:{lora_info.get('video_to_audio', 1.0):.2f}", + f"a>v:{lora_info.get('audio_to_video', 1.0):.2f}", + f"o:{lora_info.get('other', 1.0):.2f}", + ] + print( + f" {i}. {lora_name} (strength: {strength_model:.2f} | {', '.join(layer_strs)})" + ) + print(f"Target rank: {rank}") + print(f"Output prefix: {output_prefix}") + print(f"Output directory: {output_dir}") + print(f"Merge mode: {'Model-weight (accurate)' if merge_on_model else 'LoRA SVD (fast)'}") + print(f"{'=' * 60}\n") + + # Determine output path early to show user where file will be saved + output_path = self._get_output_path(output_prefix, output_dir) + print(f"Output path: {output_path}") + print() + + # Get device for computation + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Load all LoRA files and build merged state dict + lora_data_list = [] # List of (lora_sd, metadata, layer_strengths) + + for lora_info in lora_chain: + lora_name = lora_info.get("lora_name", "") + if not lora_name: + continue + + # Load LoRA file + try: + lora_path = folder_paths.get_full_path("loras", lora_name) + lora_sd, _ = self._load_lora_file(lora_path) + + layer_strengths = { + "video": lora_info.get("video", 1.0), + "video_to_audio": lora_info.get("video_to_audio", 1.0), + "audio": lora_info.get("audio", 1.0), + "audio_to_video": lora_info.get("audio_to_video", 1.0), + "other": lora_info.get("other", 1.0), + } + + lora_data_list.append( + { + "sd": lora_sd, + "name": lora_name, + "strength_model": lora_info.get("strength_model", 1.0), + "layer_strengths": layer_strengths, + } + ) + except Exception as e: + print(f"Warning: Could not load LoRA '{lora_name}': {e}") + continue + + if not lora_data_list: + raise ValueError("No valid LoRAs found to merge") + + # Get checkpoint path for model-weight merge + checkpoint_path = model_with_loras.model_options.get("checkpoint_name", "") + + # Debug: Show sample LoRA keys from first loaded LoRA + if lora_data_list: + sample_keys = list(lora_data_list[0]["sd"].keys())[:5] + print(f"Sample LoRA keys from '{lora_data_list[0]['name']}':") + for k in sample_keys: + print(f" - {k}") + + # Build merged state dict by combining all LoRA contributions + if merge_on_model and checkpoint_path: + print(f"Checkpoint path: {checkpoint_path}") + if not Path(checkpoint_path).exists(): + raise ValueError(f"Checkpoint path not found: {checkpoint_path}") + + merged_sd = merge_lora_state_dicts( + lora_data_list, rank, device, checkpoint_path=checkpoint_path + ) + else: + if merge_on_model: + print("Warning: merge_on_model=True but no checkpoint path found. Using LoRA SVD merge instead.") + merged_sd = merge_lora_state_dicts(lora_data_list, rank, device) + + # Print layer count summary + print(f"\nMerge Summary:") + print(f" Total merged: {len(merged_sd) // 3} layers") + + # Save merged LoRA + self._save_merged_lora(merged_sd, output_path, lora_chain, rank) + + print(f"\n{'=' * 60}") + print(f"Merge complete!") + print(f"Output: {output_path}") + print(f"{'=' * 60}\n") + + return (str(output_path),) + + def _load_lora_file(self, path): + """Load a LoRA file (safetensors or pickle).""" + from safetensors import safe_open as st_safe_open + + path = Path(path) + if path.suffix == ".safetensors": + with st_safe_open(path, framework="pt", device="cpu") as f: + sd = {k: f.get_tensor(k) for k in f.keys()} + metadata = f.metadata() + else: + sd = torch.load(path, map_location="cpu", weights_only=True) + if isinstance(sd, dict) and "state_dict" in sd: + sd = sd["state_dict"] + metadata = None + return sd, metadata + + def _get_output_path(self, prefix, output_dir="loras"): + """ + Determine the output path for the merged LoRA. + + Args: + prefix: Filename prefix + output_dir: Output directory - either 'output' for ComfyUI output folder, + 'loras' for ComfyUI loras folder, or a custom relative/absolute path + """ + from datetime import datetime + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"{prefix}_{timestamp}.safetensors" + + # Handle special directory names + if output_dir.lower() == "loras": + # Save to ComfyUI loras folder for auto-discovery + try: + loras_dirs = folder_paths.get_folder_paths("loras") + if loras_dirs: + target_dir = os.path.join(loras_dirs[0], prefix) + os.makedirs(target_dir, exist_ok=True) + return os.path.join(target_dir, filename) + except Exception as e: + print(f"Warning: Could not access loras directory: {e}") + # Fallback to output/loras if folder_paths doesn't work + output_base = folder_paths.get_output_directory() + target_dir = os.path.join(output_base, "loras", prefix) + elif output_dir.lower() == "output": + # Save to ComfyUI output folder using standard method + try: + full_output_folder, _, _, _ = folder_paths.get_save_image_path(prefix, None) + return os.path.join(full_output_folder, filename) + except Exception as e: + print(f"Warning: Could not get save image path: {e}") + output_base = folder_paths.get_output_directory() + target_dir = os.path.join(output_base, prefix) + else: + # Custom directory - treat as relative to ComfyUI output + output_base = folder_paths.get_output_directory() + target_dir = os.path.join(output_base, output_dir, prefix) + + os.makedirs(target_dir, exist_ok=True) + return os.path.join(target_dir, filename) + + def _save_merged_lora(self, merged_sd, output_path, lora_chain, rank): + """ + Save the merged LoRA as a safetensors file with metadata. + """ + # Build metadata + metadata = { + "ss_network_module": "diffusers", + "ss_network_dim": str(rank), + "ss_network_alpha": str(rank), + "ss_merge_source": "LTX2LoRASelect Merge", + "ss_merged_loras": json.dumps([l.get("lora_name", "unknown") for l in lora_chain]), + "format": "diffusers", + } + + # Convert tensors to cpu and ensure proper format + processed_sd = {} + for k, v in merged_sd.items(): + if isinstance(v, torch.Tensor): + processed_sd[k] = v.cpu().contiguous() + else: + processed_sd[k] = v + + # Save as safetensors + save_file(processed_sd, str(output_path), metadata=metadata) + + +# Register the nodes +NODE_CLASS_MAPPINGS = { + "LoadModelMeta": LoadModelMeta, + "LTX2LoRASelect": LTX2LoRASelect, + "LoraMergeLTX": LoraMergeLTX, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoadModelMeta": "Load Model Meta", + "LTX2LoRASelect": "LTX2 LORA Select", + "LoraMergeLTX": "Lora Merge LTX", +} diff --git a/nodes/video_nodes_extra.py b/nodes/video_nodes_extra.py index 367c030..8c2c805 100644 --- a/nodes/video_nodes_extra.py +++ b/nodes/video_nodes_extra.py @@ -5,6 +5,7 @@ import base64 import numpy as np import json + class VideoMergeABC: @classmethod def INPUT_TYPES(cls): @@ -13,8 +14,18 @@ class VideoMergeABC: "image_A": ("IMAGE",), "image_B": ("IMAGE",), "image_C": ("IMAGE",), - "FlamesOverlap": ("INT", {"default": 15, "min": 1, "max": 10000, "step": 1}), - "easing_clamp": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 2.0, "step": 0.01}), + "A_overlap": ( + "INT", + {"default": 24, "min": 0, "max": 10000, "step": 1}, + ), + "B_overlap": ( + "INT", + {"default": 24, "min": 0, "max": 10000, "step": 1}, + ), + "easing_clamp": ( + "FLOAT", + {"default": 0.5, "min": -1.0, "max": 2.0, "step": 0.01}, + ), "easing_type": (["linear", "ease_in", "ease_out", "ease_in_out"],), "resize_to_B": (["off", "on"], {"default": "off"}), }, @@ -25,26 +36,43 @@ class VideoMergeABC: }, } - RETURN_TYPES = ("IMAGE", "IMAGE") - RETURN_NAMES = ("images", "debug") + RETURN_TYPES = ("IMAGE", "IMAGE", "VIDEO_META") + RETURN_NAMES = ("images", "debug", "video_meta") FUNCTION = "merge_abc" CATEGORY = "WanVideoWrapper_QQ/video" - def merge_abc(self, image_A, image_B, image_C, FlamesOverlap, easing_clamp, easing_type, resize_to_B, **kwargs): - opacity_curve = kwargs.get('opacity_curve', '[]') + def merge_abc( + self, + image_A, + image_B, + image_C, + A_overlap, + B_overlap, + easing_clamp, + easing_type, + resize_to_B, + **kwargs, + ): + opacity_curve = kwargs.get("opacity_curve", "[]") num_b = image_B.shape[0] device = image_B.device try: opacity_y = json.loads(opacity_curve) if not isinstance(opacity_y, list) or len(opacity_y) != num_b: - y_list = self._generate_y_values(num_b, FlamesOverlap, easing_clamp, easing_type) + y_list = self._generate_y_values( + num_b, A_overlap, B_overlap, easing_clamp, easing_type + ) opacity_y = torch.tensor(y_list, dtype=torch.float32, device=device) else: # Resample the curve from JS UI to ensure consistency - opacity_y = self._resample_curve(opacity_y, num_b, FlamesOverlap, easing_clamp, easing_type, device) + opacity_y = self._resample_curve( + opacity_y, num_b, A_overlap, B_overlap, easing_clamp, easing_type, device + ) except (json.JSONDecodeError, ValueError) as e: - y_list = self._generate_y_values(num_b, FlamesOverlap, easing_clamp, easing_type) + y_list = self._generate_y_values( + num_b, A_overlap, B_overlap, easing_clamp, easing_type + ) opacity_y = torch.tensor(y_list, dtype=torch.float32, device=device) num_a, h_a, w_a, c_a = image_A.shape @@ -57,35 +85,45 @@ class VideoMergeABC: target_h = h_b target_w = w_b target_c = max(c_a, c_b, c_c) - + # Resize image_A and image_C to match image_B's dimensions - image_A = self._resize_to_target(image_A, num_a, target_h, target_w, target_c) - image_C = self._resize_to_target(image_C, num_c, target_h, target_w, target_c) + image_A = self._resize_to_target( + image_A, num_a, target_h, target_w, target_c + ) + image_C = self._resize_to_target( + image_C, num_c, target_h, target_w, target_c + ) # image_B already has the correct dimensions, but ensure channel consistency - image_B = self._resize_to_target(image_B, num_b, target_h, target_w, target_c) + image_B = self._resize_to_target( + image_B, num_b, target_h, target_w, target_c + ) else: # Original behavior: use the maximum dimensions target_h = max(h_a, h_b, h_c) target_w = max(w_a, w_b, w_c) target_c = max(c_a, c_b, c_c) - image_A = self._resize_to_target(image_A, num_a, target_h, target_w, target_c) - image_B = self._resize_to_target(image_B, num_b, target_h, target_w, target_c) - image_C = self._resize_to_target(image_C, num_c, target_h, target_w, target_c) + image_A = self._resize_to_target( + image_A, num_a, target_h, target_w, target_c + ) + image_B = self._resize_to_target( + image_B, num_b, target_h, target_w, target_c + ) + image_C = self._resize_to_target( + image_C, num_c, target_h, target_w, target_c + ) - safe_overlap = min(FlamesOverlap, num_a, num_c, num_b // 2) + safe_A_overlap = min(A_overlap, num_a, num_b // 2) + safe_B_overlap = min(B_overlap, num_c, num_b // 2) - if safe_overlap == 0: - return (image_B, image_B) # Return same image for both outputs when no overlap - - A1 = image_A[:-safe_overlap] - A2 = image_A[-safe_overlap:] - C1 = image_C[:safe_overlap] - C2 = image_C[safe_overlap:] + A1 = image_A[:-safe_A_overlap] + A2 = image_A[-safe_A_overlap:] + C1 = image_C[:safe_B_overlap] + C2 = image_C[safe_B_overlap:] B_old = torch.zeros_like(image_B) - B_old[:safe_overlap] = A2 - B_old[-safe_overlap:] = C1 + B_old[:safe_A_overlap] = A2 + B_old[-safe_B_overlap:] = C1 # Use modular blending function for actual output B_new = self._blend_images(image_B, B_old, opacity_y) @@ -95,53 +133,77 @@ class VideoMergeABC: # Create debug visualization with blue overlay on image_B debug_image = self._create_debug_image(image_A, image_B, image_C, opacity_y) - graph_data = self._create_transition_graph(num_b, FlamesOverlap, easing_clamp, easing_type) - - ui_data = {"transition_graph": [graph_data]} - return {"ui": ui_data, "result": (result, debug_image)} + graph_data = self._create_transition_graph( + num_b, A_overlap, B_overlap, easing_clamp, easing_type + ) - def _generate_y_values(self, num_frames, FlamesOverlap, easing_clamp, easing_type): + video_meta = { + "A_overlap": A_overlap, + "B_overlap": B_overlap, + "easing_clamp": easing_clamp, + "easing_type": easing_type, + "num_frames_b": num_b, + "safe_A_overlap": safe_A_overlap, + "safe_B_overlap": safe_B_overlap, + } + + ui_data = {"transition_graph": [graph_data]} + return {"ui": ui_data, "result": (result, debug_image, video_meta)} + + def _generate_y_values(self, num_frames, A_overlap, B_overlap, easing_clamp, easing_type): + """ + Generate opacity curve for blending. + """ y = [0.0] * num_frames - - safe_overlap = min(FlamesOverlap, num_frames // 2) + + safe_A_overlap = min(A_overlap, num_frames // 2) + safe_B_overlap = min(B_overlap, num_frames // 2) + # Use original remap_clamp calculation to match JS behavior remap_clamp = (easing_clamp - 0.5) * 2 - if safe_overlap > 0: - for i in range(1, safe_overlap + 1): - progress = i / safe_overlap - y_eased = self._apply_easing(progress, easing_type) - # Use original remap_clamp to allow negative values - y[i] = y_eased ** (2 ** -remap_clamp) - for i in range(safe_overlap + 1, num_frames - safe_overlap): + if safe_A_overlap > 0: + for i in range(1, safe_A_overlap + 1): + progress = i / safe_A_overlap + y_eased = self._apply_easing(progress, easing_type) + y[i] = y_eased ** (2**-remap_clamp) + + for i in range(safe_A_overlap + 1, num_frames - safe_B_overlap): y[i] = 1.0 + if safe_B_overlap > 0: inverted_easing_type = self._invert_easing_type(easing_type) - for i in range(num_frames - safe_overlap, num_frames): - progress = (i - (num_frames - safe_overlap)) / (safe_overlap - 1) if safe_overlap > 1 else 1.0 + for i in range(num_frames - safe_B_overlap, num_frames): + progress = ( + (i - (num_frames - safe_B_overlap)) / (safe_B_overlap - 1) + if safe_B_overlap > 1 + else 1.0 + ) y_eased = 1.0 - self._apply_easing(progress, inverted_easing_type) - # Use original remap_clamp to allow negative values - y[i] = y_eased ** (2 ** -remap_clamp) - + y[i] = y_eased ** (2**-remap_clamp) + if num_frames > 0: y[0] = 0.0 if num_frames > 1: - y[num_frames-1] = 0.0 + y[num_frames - 1] = 0.0 return y - def _resample_curve(self, js_curve, target_frames, FlamesOverlap, easing_clamp, easing_type, device): + def _resample_curve( + self, js_curve, target_frames, A_overlap, B_overlap, easing_clamp, easing_type, device + ): """ Resample the curve from JS UI to ensure consistency and prevent negative jumping. - + Args: js_curve: List of y values from JS UI target_frames: Number of frames to resample to - FlamesOverlap: Overlap parameter + A_overlap: Overlap parameter for A side + B_overlap: Overlap parameter for B side easing_clamp: Easing clamp parameter easing_type: Type of easing device: PyTorch device - + Returns: Resampled tensor of opacity values """ @@ -155,21 +217,28 @@ class VideoMergeABC: clamped_val = max(0.0, min(1.0, float(val))) resampled.append(clamped_val) return torch.tensor(resampled, dtype=torch.float32, device=device) - + # Otherwise, generate a new curve using the current parameters # This ensures consistency between JS and Python - y_list = self._generate_y_values(target_frames, FlamesOverlap, easing_clamp, easing_type) + y_list = self._generate_y_values( + target_frames, A_overlap, B_overlap, easing_clamp, easing_type + ) return torch.tensor(y_list, dtype=torch.float32, device=device) - def _create_transition_graph(self, num_frames, FlamesOverlap, easing_clamp, easing_type): - y = self._generate_y_values(num_frames, FlamesOverlap, easing_clamp, easing_type) + def _create_transition_graph( + self, num_frames, A_overlap, B_overlap, easing_clamp, easing_type + ): + y = self._generate_y_values( + num_frames, A_overlap, B_overlap, easing_clamp, easing_type + ) x = list(range(num_frames)) - - safe_overlap = min(FlamesOverlap, num_frames // 2) - key_frames = [0, safe_overlap, num_frames - safe_overlap, num_frames - 1] + + safe_A_overlap = min(A_overlap, num_frames // 2) + safe_B_overlap = min(B_overlap, num_frames // 2) + key_frames = [0, safe_A_overlap, num_frames - safe_B_overlap, num_frames - 1] key_values = [y[k] if 0 <= k < len(y) else 0.0 for k in key_frames] - return { "x": x, "y": y, "key_frames": key_frames, "key_values": key_values } + return {"x": x, "y": y, "key_frames": key_frames, "key_values": key_values} def _apply_easing(self, t, easing_type): if easing_type == "ease_in": @@ -189,42 +258,62 @@ class VideoMergeABC: def _resize_to_target(self, image, target_b, target_h, target_w, target_c): b, h, w, c = image.shape - + if c != target_c: # Basic channel handling - if c == 1: image = image.repeat(1, 1, 1, target_c) - else: image = image[..., :target_c] + if c == 1: + image = image.repeat(1, 1, 1, target_c) + else: + image = image[..., :target_c] if b != target_b: if b == 1: image = image.repeat(target_b, 1, 1, 1) elif b < target_b: - image = torch.cat([image, image[-1:].repeat(target_b - b, 1, 1, 1)], dim=0) + image = torch.cat( + [image, image[-1:].repeat(target_b - b, 1, 1, 1)], dim=0 + ) else: image = image[:target_b] - + if h != target_h or w != target_w: - image = image.permute(0, 3, 1, 2) - image = F.interpolate(image, size=(target_h, target_w), mode='bilinear', align_corners=False) - image = image.permute(0, 2, 3, 1) - + # Calculate scale factors for each dimension + scale_h = target_h / h if h > 0 else 1.0 + scale_w = target_w / w if w > 0 else 1.0 + + # Use the larger scale factor (scale to cover, then crop) + scale = max(scale_h, scale_w) + + new_h = int(h * scale) + new_w = int(w * scale) + + # Resize using bicubic interpolation for better quality + image = image.permute(0, 3, 1, 2) # [B, C, H, W] + image = F.interpolate(image, size=(new_h, new_w), mode="bicubic", align_corners=False) + image = image.permute(0, 2, 3, 1) # [B, H, W, C] + + # Center crop to target dimensions + start_h = (new_h - target_h) // 2 + start_w = (new_w - target_w) // 2 + image = image[:, start_h:start_h + target_h, start_w:start_w + target_w, :] + return image def _blend_images(self, base_image, overlay_image, opacity): """Modular blending function that blends two images using the given opacity. - + Args: base_image: The base image (image_B in our case) overlay_image: The image to overlay (B_old for actual output, blue for debug) opacity: Tensor of opacity values for each frame - + Returns: Blended image tensor """ opacity_4d = opacity.view(opacity.shape[0], 1, 1, 1) # Simple mix blending: base * opacity + overlay * (1 - opacity) return base_image * opacity_4d + overlay_image * (1.0 - opacity_4d) - + def _create_debug_image(self, image_A, image_B, image_C, opacity_y): """Create a debug visualization showing image_B with blue overlay based on opacity curve.""" b, h, w, c = image_B.shape @@ -235,6 +324,887 @@ class VideoMergeABC: return image_B_with_overlay +class VideoAudioMergeAB: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_A": ("IMAGE",), + "image_B": ("IMAGE",), + "FlamesOverlap": ("INT", {"default": 15, "min": 0, "max": 1000, "step": 1}), + "video_fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 120.0, "step": 0.5}), + "easing_clamp": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 2.0, "step": 0.01}), + "easing_type": (["linear", "ease_in", "ease_out", "ease_in_out"],), + "normalize_volume": ("BOOLEAN", {"default": True}), + "b_volume_gain": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "cut_mode": ("BOOLEAN", {"default": False}), + }, + "optional": { + "audio_A": ("AUDIO",), + "audio_B": ("AUDIO",), + }, + } + + RETURN_TYPES = ("IMAGE", "AUDIO", "DEBUG") + RETURN_NAMES = ("images", "audio", "debug") + FUNCTION = "merge_ab" + CATEGORY = "WanVideoWrapper_QQ/video" + + def merge_ab(self, image_A, image_B, FlamesOverlap, video_fps, easing_clamp, easing_type, normalize_volume, b_volume_gain, cut_mode, audio_A=None, audio_B=None): + device = image_A.device + num_a = image_A.shape[0] + num_b = image_B.shape[0] + + # 1. Image Pre-processing + h_a, w_a, c_a = image_A.shape[1:] + h_b, w_b, c_b = image_B.shape[1:] + target_h, target_w, target_c = max(h_a, h_b), max(w_a, w_b), max(c_a, c_b) + + image_A = self._resize_to_target(image_A, num_a, target_h, target_w, target_c) + image_B = self._resize_to_target(image_B, num_b, target_h, target_w, target_c) + + safe_overlap = min(FlamesOverlap, num_a // 2, num_b // 2) + + # 2. Video Merge Logic + if cut_mode: + # CUT MODE: image_A as background (full length), image_B overlaid at the end + # Discard first (num_b - safe_overlap) frames of B, blend remaining onto A's end + result_image = self._cut_mode_merge(image_A, image_B, safe_overlap, easing_clamp, easing_type, device) + else: + # NORMAL MODE: crossfade merge A -> B + if safe_overlap > 0: + A1 = image_A[:-safe_overlap] + A2 = image_A[-safe_overlap:] + B1 = image_B[:safe_overlap] + B2 = image_B[safe_overlap:] + + y_list = self._generate_y_values_ab(safe_overlap, easing_clamp, easing_type) + opacity_y = torch.tensor(y_list, dtype=torch.float32, device=device) + blended_img = self._blend_images(B1, A2, opacity_y) + result_image = torch.cat([A1, blended_img, B2], dim=0) + else: + result_image = torch.cat([image_A, image_B], dim=0) + + # 3. Handle audio (skip in cut_mode) + if cut_mode: + # CUT MODE: no audio processing + result_audio = {"waveform": torch.zeros((1, 2, 1), device=device), "sample_rate": 44100} + debug_info = { + "video_frames": result_image.shape[0], + "audio_samples": 0, + "overlap_frames": safe_overlap, + "sample_rate": 44100, + "mode": "cut_mode", + "note": "Cut mode - image_B overlaid at end of image_A, no audio" + } + return (result_image, result_audio, debug_info) + + # Normal mode: check if both audio inputs are missing + has_audio_a = audio_A is not None + has_audio_b = audio_B is not None + + if not has_audio_a and not has_audio_b: + # No audio provided, return video only with empty audio + result_audio = {"waveform": torch.zeros((1, 2, 1), device=device), "sample_rate": 44100} + debug_info = { + "video_frames": result_image.shape[0], + "audio_samples": 0, + "overlap_frames": safe_overlap, + "sample_rate": 44100, + "note": "No audio inputs provided - video merge only" + } + return (result_image, result_audio, debug_info) + + # 3. Audio Extraction (Fixed for LazyAudioMap) + output_sr = 44100 + if has_audio_a: + waveform_a, sr_a = self._extract_audio(audio_A) + waveform_a = self._to_std_audio(waveform_a, device) + output_sr = max(output_sr, sr_a) + else: + # Create silent audio for A with matching duration + waveform_a = torch.zeros((int(num_a * 44100 / video_fps), 2), device=device) + sr_a = 44100 + + if has_audio_b: + waveform_b, sr_b = self._extract_audio(audio_B) + waveform_b = self._to_std_audio(waveform_b, device) + output_sr = max(output_sr, sr_b) + # Apply gain to B before normalization or mixing + waveform_b = waveform_b * b_volume_gain + else: + # Create silent audio for B with matching duration + waveform_b = torch.zeros((int(num_b * 44100 / video_fps), 2), device=device) + sr_b = 44100 + + if normalize_volume and has_audio_a: + waveform_a = self._normalize_audio(waveform_a) + if normalize_volume and has_audio_b: + waveform_b = self._normalize_audio(waveform_b) + + max_channels = max(waveform_a.shape[1], waveform_b.shape[1]) + waveform_a = self._match_channels(waveform_a, max_channels) + waveform_b = self._match_channels(waveform_b, max_channels) + + # 4. Precise Audio Resampling + overlap_samples = int(safe_overlap * output_sr / video_fps) + a_part1_samples = int((num_a - safe_overlap) * output_sr / video_fps) + b_part2_samples = int((num_b - safe_overlap) * output_sr / video_fps) + + # Calculate indices based on ratios to maintain sync + audio_a_part1 = self._resample_waveform(waveform_a[:int(waveform_a.shape[0] * (num_a-safe_overlap)/num_a)], a_part1_samples) + audio_a_overlap = self._resample_waveform(waveform_a[int(waveform_a.shape[0] * (num_a-safe_overlap)/num_a):], overlap_samples) + audio_b_overlap = self._resample_waveform(waveform_b[:int(waveform_b.shape[0] * safe_overlap/num_b)], overlap_samples) + audio_b_part2 = self._resample_waveform(waveform_b[int(waveform_b.shape[0] * safe_overlap/num_b):], b_part2_samples) + + # 5. Audio Crossfade + if overlap_samples > 0: + y_list_audio = self._generate_y_values_ab(overlap_samples, easing_clamp, easing_type) + opacity_audio = torch.tensor(y_list_audio, dtype=torch.float32, device=device).unsqueeze(-1) + blended_audio = (audio_b_overlap * opacity_audio) + (audio_a_overlap * (1.0 - opacity_audio)) + result_audio_waveform = torch.cat([audio_a_part1, blended_audio, audio_b_part2], dim=0) + else: + result_audio_waveform = torch.cat([audio_a_part1, audio_b_part2], dim=0) + + result_audio = { + "waveform": result_audio_waveform.t().unsqueeze(0), + "sample_rate": output_sr, + } + + debug_info = { + "video_frames": result_image.shape[0], + "audio_samples": result_audio_waveform.shape[0], + "overlap_frames": safe_overlap, + "sample_rate": output_sr, + "has_audio_a": has_audio_a, + "has_audio_b": has_audio_b + } + + return (result_image, result_audio, debug_info) + + def _extract_audio(self, audio): + """Standardizes input that might be a dict, LazyAudioMap, or Tensor.""" + waveform = None + sample_rate = 44100 + + if hasattr(audio, "get"): # Handle LazyAudioMap + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + elif isinstance(audio, dict): # Handle standard dict + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + else: + waveform = audio + + return waveform, sample_rate + + def _to_std_audio(self, w, device): + # Final safety check before converting to tensor + if not isinstance(w, torch.Tensor): + if hasattr(w, "numpy"): w = torch.from_numpy(w()) + else: w = torch.tensor(w, dtype=torch.float32) + + w = w.to(device) + if w.ndim == 1: w = w.unsqueeze(-1) + while w.ndim > 2: w = w.squeeze(0) + if w.shape[0] < w.shape[1]: w = w.t() + return w + + def _match_channels(self, w, target_c): + if w.shape[1] < target_c: + return w.repeat(1, target_c // w.shape[1] + 1)[:, :target_c] + return w + + def _normalize_audio(self, w, target_db=-20.0): + rms = torch.sqrt(torch.mean(w**2) + 1e-9) + target_amp = 10**(target_db / 20) + return w * (target_amp / rms) + + def _resample_waveform(self, waveform, target_samples): + if target_samples <= 0: return torch.empty((0, waveform.shape[1]), device=waveform.device) + if waveform.shape[0] == 0: return torch.zeros((target_samples, waveform.shape[1]), device=waveform.device) + w = waveform.t().unsqueeze(0) + resampled = F.interpolate(w, size=int(target_samples), mode='linear', align_corners=False) + return resampled.squeeze(0).t() + + def _generate_y_values_ab(self, num_frames, easing_clamp, easing_type): + if num_frames <= 1: return [0.5] if num_frames == 1 else [] + remap_clamp = (easing_clamp - 0.5) * 2 + y = [] + for i in range(num_frames): + t = i / (num_frames - 1) + t_eased = self._apply_easing(t, easing_type) + y.append(t_eased ** (2**-remap_clamp)) + return y + + def _apply_easing(self, t, easing_type): + if easing_type == "ease_in": return t * t + if easing_type == "ease_out": return 1.0 - (1.0 - t)**2 + if easing_type == "ease_in_out": return 2*t*t if t < 0.5 else 1-2*(1-t)**2 + return t + + def _resize_to_target(self, image, target_b, target_h, target_w, target_c): + b, h, w, c = image.shape + if (h, w, c) != (target_h, target_w, target_c): + image = image.permute(0, 3, 1, 2) + image = F.interpolate(image, size=(target_h, target_w), mode="bilinear", align_corners=False) + image = image.permute(0, 2, 3, 1) + return image + + def _blend_images(self, base, overlay, opacity): + mask = opacity.view(-1, 1, 1, 1) + return base * mask + overlay * (1.0 - mask) + + def _cut_mode_merge(self, image_A, image_B, safe_overlap, easing_clamp, easing_type, device): + """ + CUT MODE: Overlay image_B at the end of image_A. + + Example: image_A=121 frames, image_B=21 frames, FlamesOverlap=5 + - Discard first (21-5)=16 frames of B + - Use last 5 frames of B blended from 0% to 100% opacity over A's last 5 frames + - Result: 121 frames total, with last frame being 100% image_B + """ + num_a = image_A.shape[0] + + # Get the last safe_overlap frames of B to overlay + b_overlay = image_B[-safe_overlap:] if safe_overlap <= image_B.shape[0] else image_B + a_background_end = image_A[-safe_overlap:] + + # Generate opacity curve from 0 to 1 (fade in B over A) + y_list = self._generate_y_values_ab(safe_overlap, easing_clamp, easing_type) + opacity_y = torch.tensor(y_list, dtype=torch.float32, device=device) + + # Blend: result = B * opacity + A * (1 - opacity) + blended_end = self._blend_images(b_overlay, a_background_end, opacity_y) + + # Concatenate: full A minus last safe_overlap frames + blended end + if safe_overlap > 0 and num_a > safe_overlap: + result_image = torch.cat([image_A[:-safe_overlap], blended_end], dim=0) + else: + result_image = blended_end + + return result_image + + +class AudioMergeABC: + """ + Merges three audio sources (A, B, C) using the same easing curve parameters as VideoMergeABC. + + The merge structure mirrors the video: + - A fades out during A→B transition zone + - B is at full volume in the middle section + - B fades out and C fades in during B→C transition zone + + Uses video_meta from VideoMergeABC to ensure synchronized transitions. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_A": ("AUDIO",), + "audio_B": ("AUDIO",), + "audio_C": ("AUDIO",), + "video_meta": ("VIDEO_META",), + "video_fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 120.0, "step": 0.5}), + "normalize_volume": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "merge_abc_audio" + CATEGORY = "WanVideoWrapper_QQ/audio" + + def merge_abc_audio( + self, + video_meta, + video_fps, + normalize_volume, + audio_A=None, + audio_B=None, + audio_C=None, + ): + # Use CPU for audio processing to avoid device mismatch issues + device = torch.device("cpu") + output_sr = 44100 + + # Extract merge settings from video_meta (from VideoMergeABC) + A_overlap = video_meta.get("A_overlap", 24) + B_overlap = video_meta.get("B_overlap", 24) + easing_clamp = video_meta.get("easing_clamp", 0.5) + easing_type = video_meta.get("easing_type", "linear") + + # Extract audio from inputs + has_audio_a = audio_A is not None + has_audio_b = audio_B is not None + has_audio_c = audio_C is not None + + if has_audio_a: + waveform_a, sr_a = self._extract_audio(audio_A) + waveform_a = self._to_std_audio(waveform_a, device) + output_sr = max(output_sr, sr_a) + else: + waveform_a = torch.zeros((int(120 * 44100 / video_fps), 2), device=device) + sr_a = 44100 + + if has_audio_b: + waveform_b, sr_b = self._extract_audio(audio_B) + waveform_b = self._to_std_audio(waveform_b, device) + output_sr = max(output_sr, sr_b) + else: + waveform_b = torch.zeros((int(120 * 44100 / video_fps), 2), device=device) + sr_b = 44100 + + if has_audio_c: + waveform_c, sr_c = self._extract_audio(audio_C) + waveform_c = self._to_std_audio(waveform_c, device) + output_sr = max(output_sr, sr_c) + else: + waveform_c = torch.zeros((int(120 * 44100 / video_fps), 2), device=device) + sr_c = 44100 + + # Normalize volume if enabled + if normalize_volume and has_audio_a: + waveform_a = self._normalize_audio(waveform_a) + if normalize_volume and has_audio_b: + waveform_b = self._normalize_audio(waveform_b) + if normalize_volume and has_audio_c: + waveform_c = self._normalize_audio(waveform_c) + + # Match channels + max_channels = max( + waveform_a.shape[1], waveform_b.shape[1], waveform_c.shape[1] + ) + waveform_a = self._match_channels(waveform_a, max_channels) + waveform_b = self._match_channels(waveform_b, max_channels) + waveform_c = self._match_channels(waveform_c, max_channels) + + # Get frame counts from audio duration (assuming they match video frames) + num_a = int(waveform_a.shape[0] * video_fps / output_sr) + num_b = int(waveform_b.shape[0] * video_fps / output_sr) + num_c = int(waveform_c.shape[0] * video_fps / output_sr) + + # Calculate safe overlaps (matching VideoMergeABC logic) + safe_A_overlap = min(A_overlap, num_a, num_b // 2) + safe_B_overlap = min(B_overlap, num_c, num_b // 2) + + # Convert frame overlaps to sample counts + a_overlap_samples = int(safe_A_overlap * output_sr / video_fps) + b_overlap_samples = int(safe_B_overlap * output_sr / video_fps) + + # Calculate part lengths in samples + a_part1_samples = int((num_a - safe_A_overlap) * output_sr / video_fps) + b_middle_samples = int((num_b - safe_A_overlap - safe_B_overlap) * output_sr / video_fps) + c_part2_samples = int((num_c - safe_B_overlap) * output_sr / video_fps) + + # Generate opacity curves for transitions + # A→B transition: B fades in (0→1), so A fades out (1→0) + if a_overlap_samples > 0: + y_ab = self._generate_y_values_ab(a_overlap_samples, easing_clamp, easing_type) + opacity_ab = torch.tensor(y_ab, dtype=torch.float32, device=device).unsqueeze(-1) + + # B→C transition: C fades in (0→1), B fades out (1→0) + if b_overlap_samples > 0: + y_bc = self._generate_y_values_ab(b_overlap_samples, easing_clamp, easing_type) + opacity_bc = torch.tensor(y_bc, dtype=torch.float32, device=device).unsqueeze(-1) + + # Split and process audio segments using ratio-based slicing for sync + # Audio A: part1 + overlap_region (for crossfade with B) + if a_part1_samples > 0: + ratio_a1_end = (num_a - safe_A_overlap) / num_a if num_a > 0 else 1.0 + idx_a1_end = int(waveform_a.shape[0] * ratio_a1_end) + audio_a_part1 = self._resample_waveform( + waveform_a[:idx_a1_end], a_part1_samples + ) + else: + audio_a_part1 = torch.zeros((0, max_channels), device=device) + + # A overlap region (will be crossfaded with B) + if a_overlap_samples > 0 and safe_A_overlap > 0: + ratio_a_fade_start = (num_a - safe_A_overlap) / num_a if num_a > 0 else 0.0 + idx_a_fade_start = int(waveform_a.shape[0] * ratio_a_fade_start) + audio_a_overlap = self._resample_waveform( + waveform_a[idx_a_fade_start:], a_overlap_samples + ) + else: + audio_a_overlap = torch.zeros((0, max_channels), device=device) + + # Audio B: overlap_region (with A) + middle + overlap_region (with C) + if safe_A_overlap > 0 and a_overlap_samples > 0: + ratio_b_fade_end = safe_A_overlap / num_b if num_b > 0 else 1.0 + idx_b_fade_end = int(waveform_b.shape[0] * ratio_b_fade_end) + audio_b_overlap_a = self._resample_waveform( + waveform_b[:idx_b_fade_end], a_overlap_samples + ) + else: + audio_b_overlap_a = torch.zeros((0, max_channels), device=device) + + if b_middle_samples > 0: + ratio_b_mid_start = safe_A_overlap / num_b if num_b > 0 else 0.0 + ratio_b_mid_end = (num_b - safe_B_overlap) / num_b if num_b > 0 else 1.0 + idx_b_mid_start = int(waveform_b.shape[0] * ratio_b_mid_start) + idx_b_mid_end = int(waveform_b.shape[0] * ratio_b_mid_end) + audio_b_middle = self._resample_waveform( + waveform_b[idx_b_mid_start:idx_b_mid_end], b_middle_samples + ) + else: + audio_b_middle = torch.zeros((0, max_channels), device=device) + + if safe_B_overlap > 0 and b_overlap_samples > 0: + ratio_b_fadeout_start = (num_b - safe_B_overlap) / num_b if num_b > 0 else 0.0 + idx_b_fadeout_start = int(waveform_b.shape[0] * ratio_b_fadeout_start) + audio_b_overlap_c = self._resample_waveform( + waveform_b[idx_b_fadeout_start:], b_overlap_samples + ) + else: + audio_b_overlap_c = torch.zeros((0, max_channels), device=device) + + # Audio C: overlap_region (with B) + part2 + if safe_B_overlap > 0 and b_overlap_samples > 0: + ratio_c_fade_end = safe_B_overlap / num_c if num_c > 0 else 1.0 + idx_c_fade_end = int(waveform_c.shape[0] * ratio_c_fade_end) + audio_c_overlap = self._resample_waveform( + waveform_c[:idx_c_fade_end], b_overlap_samples + ) + else: + audio_c_overlap = torch.zeros((0, max_channels), device=device) + + if c_part2_samples > 0: + ratio_c2_start = safe_B_overlap / num_c if num_c > 0 else 0.0 + idx_c2_start = int(waveform_c.shape[0] * ratio_c2_start) + audio_c_part2 = self._resample_waveform( + waveform_c[idx_c2_start:], c_part2_samples + ) + else: + audio_c_part2 = torch.zeros((0, max_channels), device=device) + + # Crossfade blending: gradually blend overlapping audios using the same curve as video + # A→B crossfade: B fades in (opacity 0→1), A fades out (1-opacity) + if a_overlap_samples > 0: + # blended = B * opacity + A * (1 - opacity) where opacity goes from 0 to 1 + blended_ab = (audio_b_overlap_a * opacity_ab) + (audio_a_overlap * (1.0 - opacity_ab)) + else: + blended_ab = torch.zeros((0, max_channels), device=device) + + # B→C crossfade: C fades in (opacity 0→1), B fades out (1-opacity) + if b_overlap_samples > 0: + # blended = C * opacity + B * (1 - opacity) where opacity goes from 0 to 1 + blended_bc = (audio_c_overlap * opacity_bc) + (audio_b_overlap_c * (1.0 - opacity_bc)) + else: + blended_bc = torch.zeros((0, max_channels), device=device) + + # Final concatenation: A_part1 + AB_crossfade + B_middle + BC_crossfade + C_part2 + result_waveform = torch.cat( + [ + audio_a_part1, + blended_ab, + audio_b_middle, + blended_bc, + audio_c_part2, + ], + dim=0, + ) + + # Ensure non-zero output + if result_waveform.shape[0] == 0: + result_waveform = torch.zeros((output_sr, max_channels), device=device) + + result_audio = { + "waveform": result_waveform.t().unsqueeze(0), + "sample_rate": output_sr, + } + + return (result_audio,) + + def _extract_audio(self, audio): + """Standardizes input that might be a dict, LazyAudioMap, or Tensor.""" + waveform = None + sample_rate = 44100 + + if hasattr(audio, "get"): # Handle LazyAudioMap + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + elif isinstance(audio, dict): # Handle standard dict + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + else: + waveform = audio + + return waveform, sample_rate + + def _to_std_audio(self, w, device): + """Convert audio to standard [samples, channels] tensor format.""" + if not isinstance(w, torch.Tensor): + if hasattr(w, "numpy"): + w = torch.from_numpy(w()) + else: + w = torch.tensor(w, dtype=torch.float32) + + w = w.to(device) + if w.ndim == 1: + w = w.unsqueeze(-1) + while w.ndim > 2: + w = w.squeeze(0) + # Ensure samples first, channels second + if w.shape[0] < w.shape[1]: + w = w.t() + return w + + def _match_channels(self, w, target_c): + """Match audio to target channel count.""" + if w.shape[1] < target_c: + return w.repeat(1, target_c // w.shape[1] + 1)[:, :target_c] + return w + + def _normalize_audio(self, w, target_db=-20.0): + """Normalize audio to target RMS level.""" + rms = torch.sqrt(torch.mean(w**2) + 1e-9) + target_amp = 10 ** (target_db / 20) + return w * (target_amp / rms) + + def _resample_waveform(self, waveform, target_samples): + """Resample audio waveform to target sample count using linear interpolation.""" + if target_samples <= 0: + return torch.empty((0, waveform.shape[1]), device=waveform.device) + if waveform.shape[0] == 0: + return torch.zeros( + (target_samples, waveform.shape[1]), device=waveform.device + ) + w = waveform.t().unsqueeze(0) # [1, channels, samples] + resampled = F.interpolate( + w, size=int(target_samples), mode="linear", align_corners=False + ) + return resampled.squeeze(0).t() # [samples, channels] + + def _generate_y_values_ab(self, num_frames, easing_clamp, easing_type): + """Generate fade curve from 0 to 1 for AB-style transition.""" + if num_frames <= 1: + return [0.5] if num_frames == 1 else [] + remap_clamp = (easing_clamp - 0.5) * 2 + y = [] + for i in range(num_frames): + t = i / (num_frames - 1) + t_eased = self._apply_easing(t, easing_type) + y.append(t_eased ** (2**-remap_clamp)) + return y + + def _apply_easing(self, t, easing_type): + """Apply easing function to normalized time value.""" + if easing_type == "ease_in": + return t * t + if easing_type == "ease_out": + return 1.0 - (1.0 - t) ** 2 + if easing_type == "ease_in_out": + return 2 * t * t if t < 0.5 else 1 - 2 * (1 - t) ** 2 + return t + + +class VideoPrepAB: + """ + Prepares a video sequence by combining end frames from image_A, start frames from image_B, + and inserting empty colored frames in the middle. + + Uses image_A's size as reference - scales image_B to match on one dimension + and crops the larger one. + + Optionally handles audio the same way: end of audio_A + silent mid + start of audio_B. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_A": ("IMAGE",), + "image_B": ("IMAGE",), + "mid_frames_color": ("FLOAT", {"default": 50, "min": 0, "max": 100, "step": 1}), + "end_A_frames": ("INT", {"default": 24, "min": 0, "max": 10000, "step": 1}), + "start_B_frames": ("INT", {"default": 24, "min": 0, "max": 10000, "step": 1}), + "total_frames": ("INT", {"default": 121, "min": 1, "max": 993, "step": 8}), + "video_fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 120.0, "step": 0.5}), + "audio_blend_frames": ("INT", {"default": 4, "min": 0, "max": 1000, "step": 1}), + }, + "optional": { + "audio_A": ("AUDIO",), + "audio_B": ("AUDIO",), + }, + } + + RETURN_TYPES = ("IMAGE", "AUDIO", "IMAGE") + RETURN_NAMES = ("Images", "audio", "audio_mask") + FUNCTION = "video_prep_ab" + CATEGORY = "WanVideoWrapper_QQ/video" + + def video_prep_ab(self, image_A, image_B, mid_frames_color, end_A_frames, start_B_frames, total_frames, video_fps, audio_blend_frames, audio_A=None, audio_B=None): + device = image_A.device + + # Get dimensions + num_a, h_a, w_a, c_a = image_A.shape + num_b, h_b, w_b, c_b = image_B.shape + + # Use image_A's dimensions as target + target_h, target_w, target_c = h_a, w_a, max(c_a, c_b) + + # Resize/scale image_B to match image_A's size (scale on matching side, crop biggest) + image_B_resized = self._resize_to_match(image_B, num_b, target_h, target_w, target_c) + + # Ensure image_A has correct channels + if c_a != target_c: + if c_a == 1: + image_A = image_A.repeat(1, 1, 1, target_c) + else: + image_A = image_A[..., :target_c] + + # Calculate mid frames count + mid_frames_count = total_frames - end_A_frames - start_B_frames + + # Clamp frame counts to available frames + end_A_frames = min(end_A_frames, num_a) + start_B_frames = min(start_B_frames, image_B_resized.shape[0]) + + # If mid_frames_count is negative, we need to reduce total or the overlap frames + if mid_frames_count < 0: + # Reduce end_A and start_B proportionally + excess = -mid_frames_count + end_A_frames = max(0, end_A_frames - excess // 2) + start_B_frames = max(0, start_B_frames - (excess + excess % 2) // 2) + mid_frames_count = total_frames - end_A_frames - start_B_frames + + # Get frames from image_A (last end_A_frames) + if end_A_frames > 0: + frames_a = image_A[-end_A_frames:] if end_A_frames <= num_a else image_A + else: + frames_a = torch.zeros((0, target_h, target_w, target_c), device=device) + + # Get frames from image_B (first start_B_frames) + if start_B_frames > 0: + frames_b = image_B_resized[:start_B_frames] + else: + frames_b = torch.zeros((0, target_h, target_w, target_c), device=device) + + # Create mid frames with color based on mid_frames_color (0-100 -> 0.0-1.0) + color_value = mid_frames_color / 100.0 + if mid_frames_count > 0: + mid_frames = torch.full( + (mid_frames_count, target_h, target_w, target_c), + color_value, + device=device + ) + else: + mid_frames = torch.zeros((0, target_h, target_w, target_c), device=device) + + # Concatenate: A_end + mid_empty + B_start + result = torch.cat([frames_a, mid_frames, frames_b], dim=0) + + # === Audio Processing (mirrors video structure) === + output_sr = 44100 + has_audio_a = audio_A is not None + has_audio_b = audio_B is not None + + if has_audio_a or has_audio_b: + # Extract audio from inputs + if has_audio_a: + waveform_a, sr_a = self._extract_audio(audio_A) + waveform_a = self._to_std_audio(waveform_a, device) + output_sr = max(output_sr, sr_a) + else: + waveform_a = torch.zeros((int(total_frames * 44100 / video_fps), 2), device=device) + + if has_audio_b: + waveform_b, sr_b = self._extract_audio(audio_B) + waveform_b = self._to_std_audio(waveform_b, device) + output_sr = max(output_sr, sr_b) + else: + waveform_b = torch.zeros((int(total_frames * 44100 / video_fps), 2), device=device) + + # Match channels + max_channels = max(waveform_a.shape[1], waveform_b.shape[1]) + waveform_a = self._match_channels(waveform_a, max_channels) + waveform_b = self._match_channels(waveform_b, max_channels) + + # Calculate sample counts from frame counts + end_A_samples = int(end_A_frames * output_sr / video_fps) + start_B_samples = int(start_B_frames * output_sr / video_fps) + mid_samples = int(mid_frames_count * output_sr / video_fps) + + # Get audio segments (mirroring video logic) + num_a = waveform_a.shape[0] + num_b = waveform_b.shape[0] + + # Audio A: last end_A_samples + if end_A_samples > 0 and num_a > 0: + actual_end_A = min(end_A_samples, num_a) + audio_a = self._resample_waveform(waveform_a[-actual_end_A:], end_A_samples) + else: + audio_a = torch.zeros((0, max_channels), device=device) + + # Audio B: first start_B_samples + if start_B_samples > 0 and num_b > 0: + actual_start_B = min(start_B_samples, num_b) + audio_b = self._resample_waveform(waveform_b[:actual_start_B], start_B_samples) + else: + audio_b = torch.zeros((0, max_channels), device=device) + + # Mid: silent samples + if mid_samples > 0: + audio_mid = torch.zeros((mid_samples, max_channels), device=device) + else: + audio_mid = torch.zeros((0, max_channels), device=device) + + # Concatenate: A_end + mid_silent + B_start + result_waveform = torch.cat([audio_a, audio_mid, audio_b], dim=0) + else: + # No audio inputs - create silent output matching video duration + total_samples = int(total_frames * output_sr / video_fps) + result_waveform = torch.zeros((total_samples, 2), device=device) + + result_audio = { + "waveform": result_waveform.t().unsqueeze(0), + "sample_rate": output_sr, + } + + # === Audio Mask Generation === + # Creates a mask mirroring the video structure: + # - White (1.0) for end_A_frames region (A's portion) + # - Black (0.0) for mid frames region + # - White (1.0) for start_B_frames region (B's portion) + # With linear transitions using audio_blend_frames on both sides + + total_samples = int(total_frames * output_sr / video_fps) + end_A_samples = int(end_A_frames * output_sr / video_fps) + start_B_samples = int(start_B_frames * output_sr / video_fps) + blend_samples = max(1, int(audio_blend_frames * output_sr / video_fps)) + + # Create mask: black (0.0) everywhere initially + audio_mask = torch.zeros((total_samples,), dtype=torch.float32, device=device) + + # White region at start (A's portion): samples 0 to end_A_samples + # With fade-out transition over blend_samples + if end_A_samples > 0: + white_end = max(0, end_A_samples - blend_samples) + audio_mask[:white_end] = 1.0 + # Linear fade from 1 to 0 + fade_start = white_end + fade_end = min(end_A_samples, total_samples) + if fade_end > fade_start: + fade_len = fade_end - fade_start + for i in range(fade_start, fade_end): + progress = (i - fade_start) / max(1, fade_len - 1) + audio_mask[i] = 1.0 - progress + + # White region at end (B's portion): samples (total - start_B_samples) to total + # With fade-in transition over blend_samples + if start_B_samples > 0: + b_region_start = max(0, total_samples - start_B_samples) + white_start = min(b_region_start + blend_samples, total_samples) + audio_mask[white_start:] = 1.0 + # Linear fade from 0 to 1 + if b_region_start < white_start: + fade_len = white_start - b_region_start + for i in range(b_region_start, white_start): + progress = (i - b_region_start) / max(1, fade_len - 1) + audio_mask[i] = progress + + # Reshape to match video dimensions [total_frames, height, width, channels] for IMAGE output + # Create one frame per video frame, each showing the mask value at that point in time + # Each frame is target_h x target_w x 3 (RGB grayscale), filled with the corresponding mask value + audio_mask_per_frame = torch.zeros((total_frames, target_h, target_w, 3), dtype=torch.float32, device=device) + for f in range(total_frames): + sample_idx = min(int(f * total_samples / total_frames), total_samples - 1) + mask_value = audio_mask[sample_idx] + audio_mask_per_frame[f] = mask_value # Broadcast scalar to entire frame (all 3 channels) + + return (result, result_audio, audio_mask_per_frame) + + def _resize_to_match(self, image, num_frames, target_h, target_w, target_c): + """ + Resize image to match target dimensions by scaling on one axis and cropping the other. + Scales to fit within both dimensions while maintaining aspect ratio, + then centers crops/pads as needed. + """ + b, h, w, c = image.shape + + # Calculate scale factors for each dimension + scale_h = target_h / h if h > 0 else 1.0 + scale_w = target_w / w if w > 0 else 1.0 + + # Use the larger scale factor (scale to cover, then crop) + scale = max(scale_h, scale_w) + + new_h = int(h * scale) + new_w = int(w * scale) + + # Resize using bicubic interpolation for better quality + image = image.permute(0, 3, 1, 2) # [B, C, H, W] + image = F.interpolate(image, size=(new_h, new_w), mode="bicubic", align_corners=False) + image = image.permute(0, 2, 3, 1) # [B, H, W, C] + + # Center crop to target dimensions + start_h = (new_h - target_h) // 2 + start_w = (new_w - target_w) // 2 + image = image[:, start_h:start_h + target_h, start_w:start_w + target_w, :] + + # Handle channels + if c != target_c: + if c == 1: + image = image.repeat(1, 1, 1, target_c) + else: + image = image[..., :target_c] + + return image + + def _extract_audio(self, audio): + """Standardizes input that might be a dict, LazyAudioMap, or Tensor.""" + waveform = None + sample_rate = 44100 + + if hasattr(audio, "get"): # Handle LazyAudioMap + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + elif isinstance(audio, dict): # Handle standard dict + waveform = audio.get("waveform") + sample_rate = audio.get("sample_rate", 44100) + else: + waveform = audio + + return waveform, sample_rate + + def _to_std_audio(self, w, device): + """Convert audio to standard [samples, channels] tensor format.""" + if not isinstance(w, torch.Tensor): + if hasattr(w, "numpy"): + w = torch.from_numpy(w()) + else: + w = torch.tensor(w, dtype=torch.float32) + + w = w.to(device) + if w.ndim == 1: + w = w.unsqueeze(-1) + while w.ndim > 2: + w = w.squeeze(0) + # Ensure samples first, channels second + if w.shape[0] < w.shape[1]: + w = w.t() + return w + + def _match_channels(self, w, target_c): + """Match audio to target channel count.""" + if w.shape[1] < target_c: + return w.repeat(1, target_c // w.shape[1] + 1)[:, :target_c] + return w + + def _resample_waveform(self, waveform, target_samples): + """Resample audio waveform to target sample count using linear interpolation.""" + if target_samples <= 0: + return torch.empty((0, waveform.shape[1]), device=waveform.device) + if waveform.shape[0] == 0: + return torch.zeros((target_samples, waveform.shape[1]), device=waveform.device) + w = waveform.t().unsqueeze(0) # [1, channels, samples] + resampled = F.interpolate( + w, size=int(target_samples), mode="linear", align_corners=False + ) + return resampled.squeeze(0).t() # [samples, channels] + class VideoRGBAnalysis: """ @@ -272,13 +1242,7 @@ class VideoRGBAnalysis: # Handle edge cases if b == 0: - empty_graph = { - "x": [], - "y_r": [], - "y_g": [], - "y_b": [], - "y_brightness": [] - } + empty_graph = {"x": [], "y_r": [], "y_g": [], "y_b": [], "y_brightness": []} ui_data = {"rgb_analysis_graph": [empty_graph]} return {"ui": ui_data} @@ -288,16 +1252,18 @@ class VideoRGBAnalysis: frame_means = torch.mean(image.float(), dim=[1, 2]) # [B, 3] # Extract individual channels - red_values = frame_means[:, 0].cpu().tolist() # R channel - green_values = frame_means[:, 1].cpu().tolist() # G channel - blue_values = frame_means[:, 2].cpu().tolist() # B channel + red_values = frame_means[:, 0].cpu().tolist() # R channel + green_values = frame_means[:, 1].cpu().tolist() # G channel + blue_values = frame_means[:, 2].cpu().tolist() # B channel # Calculate brightness using Rec. 601 luminance formula # brightness = 0.299 * R + 0.587 * G + 0.114 * B frame_means_cpu = frame_means.cpu() - brightness_values = (0.299 * frame_means_cpu[:, 0] + - 0.587 * frame_means_cpu[:, 1] + - 0.114 * frame_means_cpu[:, 2]).tolist() + brightness_values = ( + 0.299 * frame_means_cpu[:, 0] + + 0.587 * frame_means_cpu[:, 1] + + 0.114 * frame_means_cpu[:, 2] + ).tolist() # Build graph data x = list(range(b)) @@ -331,20 +1297,72 @@ class VideoRGBAnalysis: "y_g": y_g, "y_b": y_b, "y_brightness": brightness_values, - "display_mode": display_mode + "display_mode": display_mode, } # Return in UI format (no output, just graph) ui_data = {"rgb_analysis_graph": [graph_data]} return {"ui": ui_data} +class VideoLoopMove: + """ + Moves the last N frames of a video loop to the beginning. + Useful for adjusting loop start points without re-rendering. + + Example: 100-frame video with num_frames=16 + - Takes last 16 frames (frames 84-99) + - Puts them at the start + - Result order: [84, 85, ..., 99, 0, 1, ..., 83] + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "num_frames": ("INT", {"default": 16, "min": 0, "max": 10000, "step": 1}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "loop_move" + CATEGORY = "WanVideoWrapper_QQ/video" + + def loop_move(self, image, num_frames): + total_frames = image.shape[0] + + # Handle edge cases + if num_frames <= 0: + return (image,) + if num_frames >= total_frames: + # If num_frames >= total, just reverse or return as-is + # Let's wrap it: effectively no change needed + return (image,) + + # Split: last num_frames go to front, rest follow + moved_part = image[-num_frames:] # Last N frames + remaining_part = image[:-num_frames] # First (total - N) frames + + result = torch.cat([moved_part, remaining_part], dim=0) + + return (result,) + NODE_CLASS_MAPPINGS = { "VideoMergeABC": VideoMergeABC, + "VideoAudioMergeAB": VideoAudioMergeAB, + "AudioMergeABC": AudioMergeABC, + "VideoPrepAB": VideoPrepAB, "VideoRGBAnalysis": VideoRGBAnalysis, + "VideoLoopMove": VideoLoopMove, } NODE_DISPLAY_NAME_MAPPINGS = { "VideoMergeABC": "Video Merge ABC", + "VideoAudioMergeAB": "VideoAudio Merge AB", + "AudioMergeABC": "Audio Merge ABC", + "VideoPrepAB": "Video Prep AB", "VideoRGBAnalysis": "Video RGB Analysis", + "VideoLoopMove": "Video Loop Move", } diff --git a/utility/ltx/__init__.py b/utility/ltx/__init__.py new file mode 100644 index 0000000..7fb247c --- /dev/null +++ b/utility/ltx/__init__.py @@ -0,0 +1,48 @@ +"""LTX utilities for LoRA operations.""" + +from .core.ops import ( + merge_linear_block, + decompose_to_rank, + resize_lora_block, + apply_dare, + apply_magprune, + apply_lambda_scaling, + apply_ties, + decoupled_magnitude_direction_merge +) + +from .core.merger import merge_lora_state_dicts, count_layer_types + +from .utils.io import load_lora, save_lora +from .utils.naming import ( + LORA_DOWN_UP_FORMATS, + is_audio_layer, + is_video_ff_layer, + is_video_only_layer, + normalize_lora_key, + build_block_map +) + +from .utils.progress import ProgressBar + +__all__ = [ + 'merge_linear_block', + 'decompose_to_rank', + 'resize_lora_block', + 'apply_dare', + 'apply_magprune', + 'apply_lambda_scaling', + 'apply_ties', + 'decoupled_magnitude_direction_merge', + 'merge_lora_state_dicts', + 'count_layer_types', + 'load_lora', + 'save_lora', + 'LORA_DOWN_UP_FORMATS', + 'is_audio_layer', + 'is_video_ff_layer', + 'is_video_only_layer', + 'normalize_lora_key', + 'build_block_map', + 'ProgressBar' +] diff --git a/utility/ltx/core/__init__.py b/utility/ltx/core/__init__.py new file mode 100644 index 0000000..47de5f2 --- /dev/null +++ b/utility/ltx/core/__init__.py @@ -0,0 +1,23 @@ +"""LTX core operations for LoRA merging.""" + +from .ops import ( + merge_linear_block, + decompose_to_rank, + resize_lora_block, + apply_dare, + apply_magprune, + apply_lambda_scaling, + apply_ties, + decoupled_magnitude_direction_merge +) + +__all__ = [ + 'merge_linear_block', + 'decompose_to_rank', + 'resize_lora_block', + 'apply_dare', + 'apply_magprune', + 'apply_lambda_scaling', + 'apply_ties', + 'decoupled_magnitude_direction_merge' +] diff --git a/utility/ltx/core/merger.py b/utility/ltx/core/merger.py new file mode 100644 index 0000000..d63fc09 --- /dev/null +++ b/utility/ltx/core/merger.py @@ -0,0 +1,436 @@ +"""LoRA merging functionality for LTX models.""" + +import torch +from ..utils.progress import ProgressBar +from ..utils.naming import LORA_DOWN_UP_FORMATS + +DEFAULT_PROCESS_DTYPE = torch.bfloat16 + + +def get_layer_strength(block_name, layer_strengths): + """ + Get the appropriate strength multiplier for a given layer name. + + Args: + block_name: Name of the layer block + layer_strengths: Dictionary of layer-specific strengths + + Returns: + The strength multiplier for this layer type + """ + key_str = str(block_name).lower() + + # Video to audio cross-attention (most specific) + if "video_to_audio_attn" in key_str: + return layer_strengths["video_to_audio"] + # Audio to video cross-attention (new in updated LTX model) + elif "audio_to_video_attn" in key_str: + return layer_strengths["audio_to_video"] + # Audio attention layers + elif ".audio_attn" in key_str or "audio_ff.net" in key_str: + return layer_strengths["audio"] + # Video attention and FF layers (most general) + elif ".attn" in key_str or ".ff.net" in key_str: + return layer_strengths["video"] + # Everything else + else: + return layer_strengths["other"] + + +def merge_on_model_weights(all_blocks, lora_data_list, target_rank, device, checkpoint_path, show_progress): + """ + Merge LoRAs on model weights (more accurate than pure LoRA SVD merge). + + Algorithm: + 1. Load base model weight for current block only (block swapping to save VRAM) + 2. Apply LoRAs to base weight: W_modified = W_base + sum(strength_i * (up_i @ down_i)) + 3. Extract delta: delta = W_modified - W_base + 4. Apply SVD to delta to get new up/down matrices + 5. Move result to CPU, free memory, repeat for next block + + Args: + all_blocks: Dictionary of block_name -> list of LoRA data + lora_data_list: List of LoRA data dictionaries + target_rank: Target rank for merged LoRA + device: Device to use for computation + checkpoint_path: Path to base model checkpoint + show_progress: Whether to show progress bar + + Returns: + Merged state dictionary + """ + from safetensors.torch import safe_open + + if show_progress: + print(f"Loading base model from: {checkpoint_path}") + + # Open checkpoint once, but load weights one at a time (block swapping) + with safe_open(checkpoint_path, framework="pt", device="cpu") as f: + checkpoint_keys = list(f.keys()) + + if show_progress: + print(f"Checkpoint has {len(checkpoint_keys)} total keys") + + merged_sd = {} + total_blocks = len(all_blocks) + + if show_progress: + print(f"\nMerging {total_blocks} layers on model weights (block-swapping mode)...") + progress_bar = ProgressBar(total_blocks, "Model-weight merge") + progress_bar.start() + + processed_count = 0 + skipped_count = 0 + + for block_idx, (block_name, blocks) in enumerate(all_blocks.items()): + if not blocks: + continue + + is_conv2d = blocks[0]["is_conv2d"] + down_name = blocks[0]["down_name"] + up_name = blocks[0]["up_name"] + + # Map LoRA keys to model keys + model_key = ( + down_name + .replace(".lora_down", "") + .replace(".lora_up", "") + .replace(".lora_A", "") + .replace(".lora_B", "") + .replace(".down", "") + .replace(".up", "") + ) + + # Handle checkpoint prefix mismatch + if not model_key.startswith("model."): + model_key = "model." + model_key + + if model_key not in f.keys(): + skipped_count += 1 + if show_progress and skipped_count <= 3: + print(f" Warning: Model key not found for block '{block_name}'") + print(f" Expected checkpoint key: {model_key}") + if show_progress: + progress_bar.update(1) + continue + + # Load ONLY this block's weight from checkpoint (block swapping!) + base_weight = f.get_tensor(model_key).to(DEFAULT_PROCESS_DTYPE).to(device) + + # Start with base weight + modified_weight = base_weight.clone() + + # Apply each LoRA: W += strength * (up @ down) + for block in blocks: + strength = block["strength"] + if abs(strength) < 1e-6: + continue + + down = block["down"].to(device).to(DEFAULT_PROCESS_DTYPE) + up = block["up"].to(device).to(DEFAULT_PROCESS_DTYPE) + + if is_conv2d: + down_flat = down.reshape(down.shape[0], -1) + up_flat = up.reshape(up.shape[0], -1) + delta = (up_flat @ down_flat) * strength + delta = delta.reshape(base_weight.shape[0], base_weight.shape[1], + base_weight.shape[2], base_weight.shape[3]) + else: + delta = (up @ down) * strength + + modified_weight += delta + + # Extract delta + delta_weight = modified_weight - base_weight + + # Clear memory before SVD + del base_weight, modified_weight + if str(device) == "cuda": + torch.cuda.empty_cache() + + if delta_weight.norm() < 1e-8: + # Exception: preserve to_gate_logits layers even with low norm + is_gate_logits = "to_gate_logits" in block_name.lower() + if not is_gate_logits: + del delta_weight + if show_progress: + progress_bar.update(1) + continue + + # Apply SVD to factorize at target rank (SVD requires float32) + # For to_gate_logits, use higher effective rank to preserve information + is_gate_logits = "to_gate_logits" in block_name.lower() + effective_rank = target_rank if not is_gate_logits else max(target_rank * 4, 256) + + if is_conv2d: + out_c = delta_weight.shape[0] + in_c = delta_weight.shape[1] + kh, kw = delta_weight.shape[2], delta_weight.shape[3] + + delta_flat = delta_weight.reshape(out_c, -1) + U, S, Vh = torch.linalg.svd(delta_flat.float()) + actual_rank = min(effective_rank, len(S), U.shape[1], Vh.shape[0]) + + U = U[:, :actual_rank] @ torch.diag(S[:actual_rank]) + Vh = Vh[:actual_rank, :] + + new_up = U.reshape(out_c, actual_rank, 1, 1).to(DEFAULT_PROCESS_DTYPE).cpu() + new_down = Vh.reshape(actual_rank, in_c, kh, kw).to(DEFAULT_PROCESS_DTYPE).cpu() + else: + U, S, Vh = torch.linalg.svd(delta_weight.float()) + actual_rank = min(effective_rank, len(S), U.shape[1], Vh.shape[0]) + + U = U[:, :actual_rank] @ torch.diag(S[:actual_rank]) + Vh = Vh[:actual_rank, :] + + new_up = U.to(DEFAULT_PROCESS_DTYPE).cpu() + new_down = Vh.to(DEFAULT_PROCESS_DTYPE).cpu() + + # Clear memory and save to CPU + del delta_weight, U, S, Vh + if str(device) == "cuda": + torch.cuda.empty_cache() + + merged_sd[down_name] = new_down + merged_sd[up_name] = new_up + merged_sd[f"{block_name}.alpha"] = torch.tensor(actual_rank, dtype=torch.int) + + processed_count += 1 + + if show_progress: + progress_bar.update(1) + + if show_progress: + progress_bar.finish() + print(f" Processed: {processed_count} layers, Skipped (not in model): {skipped_count} layers") + + return merged_sd + + +def merge_lora_state_dicts(lora_data_list, target_rank, device, show_progress=True, checkpoint_path=None): + """ + Merge multiple LoRA state dicts into one. + + For each block (down/up pair), we: + 1. Compute the full weight matrix: W = up @ down for each LoRA + 2. Apply layer-specific and overall strength multipliers + 3. Sum all contributions + 4. Apply SVD to factorize at target rank + + Args: + lora_data_list: List of dictionaries with 'sd', 'strength_model', and 'layer_strengths' + target_rank: Target rank for the merged LoRA + device: Device to use for computation + show_progress: Whether to show progress bar + + Returns: + Merged state dictionary + """ + # Build a map of block_name -> list of (down, up, alpha) from each LoRA + all_blocks = {} + + for lora_data in lora_data_list: + sd = lora_data["sd"] + strength_model = lora_data["strength_model"] + layer_strengths = lora_data["layer_strengths"] + + for key in sd.keys(): + if key.endswith(".alpha"): + continue + + key_parts = key.split(".") + for fmt_down, fmt_up in LORA_DOWN_UP_FORMATS: + if len(key_parts) >= 2 and key_parts[-2] == fmt_down: + block_name = ".".join(key_parts[:-2]) + weight_suffix = ( + "." + key_parts[-1] + if key_parts[-1] in ["weight", "bias"] + else "" + ) + up_key = f"{block_name}.{fmt_up}{weight_suffix}" + + if up_key in sd: + if block_name not in all_blocks: + all_blocks[block_name] = [] + + # Get layer-specific strength + layer_str = get_layer_strength(block_name, layer_strengths) + combined_strength = strength_model * layer_str + + all_blocks[block_name].append({ + "down": sd[key], + "up": sd[up_key], + "alpha": sd.get(f"{block_name}.alpha"), + "strength": combined_strength, + "is_conv2d": len(sd[key].shape) == 4, + "down_name": key, + "up_name": up_key, + }) + break + + # Model-weight merge if checkpoint_path is provided + if checkpoint_path: + return merge_on_model_weights( + all_blocks, lora_data_list, target_rank, device, checkpoint_path, show_progress + ) + + # Initialize merged state dict + merged_sd = {} + + # Separate single-LoRA blocks from multi-LoRA blocks + single_lora_blocks = [] + multi_lora_blocks = [] + + for block_name, blocks in all_blocks.items(): + if len(blocks) == 1: + single_lora_blocks.append((block_name, blocks[0])) + else: + multi_lora_blocks.append((block_name, blocks)) + + # Process single-LoRA blocks (no SVD needed, just scale) + for block_name, block in single_lora_blocks: + strength = block["strength"] + is_conv2d = block["is_conv2d"] + down_name = block["down_name"] + up_name = block["up_name"] + alpha = block["alpha"] + + # For single LoRA, just scale the down matrix (faster than SVD) + merged_sd[down_name] = block["down"] * strength + merged_sd[up_name] = block["up"] + if alpha is not None: + merged_sd[f"{block_name}.alpha"] = alpha * strength + + # Process multi-LoRA blocks with SVD + if multi_lora_blocks: + total_blocks = len(multi_lora_blocks) + + if show_progress: + print(f"\nMerging {total_blocks} multi-LoRA layers...") + progress_bar = ProgressBar(total_blocks, "Merging") + progress_bar.start() + + processed_count = 0 + + for block_name, blocks in multi_lora_blocks: + if not blocks: + continue + + processed_count += 1 + is_conv2d = blocks[0]["is_conv2d"] + down_name = blocks[0]["down_name"] + up_name = blocks[0]["up_name"] + + # Compute merged weight matrix by summing all contributions + merged_weight = None + + for block in blocks: + if abs(block["strength"]) < 1e-6: + continue + + down = block["down"].to(device).to(DEFAULT_PROCESS_DTYPE) + up = block["up"].to(device).to(DEFAULT_PROCESS_DTYPE) + strength = block["strength"] + + if is_conv2d: + # Flatten conv for matrix multiplication + up_flat = up.reshape(up.shape[0], -1) + down_flat = down.reshape(down.shape[0], -1) + weight_contrib = (up_flat @ down_flat) * strength + else: + weight_contrib = (up @ down) * strength + + if merged_weight is None: + merged_weight = weight_contrib + else: + merged_weight = merged_weight + weight_contrib + + if merged_weight is None or merged_weight.norm() < 1e-8: + # Exception: preserve to_gate_logits layers even with low norm + # (they may have important gating information) + is_gate_logits = "to_gate_logits" in block_name.lower() + if not is_gate_logits: + if show_progress: + progress_bar.update(1) + continue + + # Apply SVD to factorize at target rank + # For to_gate_logits, use higher effective rank to preserve information + is_gate_logits = "to_gate_logits" in block_name.lower() + effective_rank = target_rank if not is_gate_logits else max(target_rank * 4, 256) + + if is_conv2d: + out_c = merged_weight.shape[0] + in_c = merged_weight.shape[1] + # Get kernel size from original down + _, _, kh, kw = blocks[0]["down"].shape + + U, S, Vh = torch.linalg.svd(merged_weight.float()) + actual_rank = min(effective_rank, len(S), U.shape[1], Vh.shape[0]) + + # Truncate and apply singular values to U + U = U[:, :actual_rank] @ torch.diag(S[:actual_rank]) + Vh = Vh[:actual_rank, :] + + new_up = U.reshape(out_c, actual_rank, 1, 1).to(DEFAULT_PROCESS_DTYPE).cpu() + new_down = Vh.reshape(actual_rank, in_c, kh, kw).to(DEFAULT_PROCESS_DTYPE).cpu() + else: + U, S, Vh = torch.linalg.svd(merged_weight.float()) + actual_rank = min(effective_rank, len(S), U.shape[1], Vh.shape[0]) + + # Truncate and apply singular values to U + U = U[:, :actual_rank] @ torch.diag(S[:actual_rank]) + Vh = Vh[:actual_rank, :] + + new_up = U.to(DEFAULT_PROCESS_DTYPE).cpu() + new_down = Vh.to(DEFAULT_PROCESS_DTYPE).cpu() + + # Store in merged state dict + merged_sd[down_name] = new_down + merged_sd[up_name] = new_up + merged_sd[f"{block_name}.alpha"] = torch.tensor(actual_rank, dtype=torch.int) + + if show_progress: + progress_bar.update(1) + + if show_progress: + progress_bar.finish() + + if show_progress: + progress_bar.finish() + + return merged_sd + + +def count_layer_types(all_blocks): + """ + Count the number of layers by type. + + Args: + all_blocks: Dictionary of block_name -> list of blocks + + Returns: + Dictionary with layer type counts + """ + layer_counts = { + "video": 0, + "audio": 0, + "video_to_audio": 0, + "audio_to_video": 0, + "other": 0, + } + + for block_name in all_blocks.keys(): + key_lower = block_name.lower() + if "video_to_audio_attn" in key_lower: + layer_counts["video_to_audio"] += 1 + elif "audio_to_video_attn" in key_lower: + layer_counts["audio_to_video"] += 1 + elif ".audio_attn" in key_lower or "audio_ff.net" in key_lower: + layer_counts["audio"] += 1 + elif ".attn" in key_lower or ".ff.net" in key_lower: + layer_counts["video"] += 1 + else: + layer_counts["other"] += 1 + + return layer_counts diff --git a/utility/ltx/core/ops.py b/utility/ltx/core/ops.py new file mode 100644 index 0000000..4d7a86f --- /dev/null +++ b/utility/ltx/core/ops.py @@ -0,0 +1,552 @@ +import torch + +DEFAULT_PROCESS_DTYPE = torch.bfloat16 + +def merge_linear_block(block_loras, device, process_dtype=DEFAULT_PROCESS_DTYPE): + """Core math for merging Linear LoRA layers.""" + first_lora = block_loras[0] + in_rank, in_size = first_lora["down"].shape + out_size, out_rank = first_lora["up"].shape + + merged_weight = torch.zeros(out_size, in_size, dtype=process_dtype, device=device) + for lora_data in block_loras: + strength = lora_data["strength"] + if abs(strength) < 1e-6: continue + + lora_down = lora_data["down"].to(process_dtype).to(device) + lora_up = lora_data["up"].to(process_dtype).to(device) + merged_weight += (lora_up @ lora_down) * strength + + return merged_weight + +def decompose_to_rank(merged_weight, target_rank): + """Performs SVD and truncates to target_rank.""" + U, S, Vh = torch.linalg.svd(merged_weight.float()) + actual_rank = min(target_rank, len(S)) + + U = U[:, :actual_rank] @ torch.diag(S[:actual_rank]) + Vh = Vh[:actual_rank, :] + return U, Vh, actual_rank + +def resize_lora_block(lora_down, lora_up, new_rank, alpha=None, device="cpu", process_dtype=DEFAULT_PROCESS_DTYPE, create_alpha=True): + """Resize a single LoRA block to new rank using SVD. + + Args: + lora_down: Down projection weight tensor + lora_up: Up projection weight tensor + new_rank: Target rank for resize + alpha: Original alpha value (tensor or scalar) + device: Device to use for computation + process_dtype: Data type for computation + create_alpha: If True, calculate and return new_alpha. If False, return None. + + Returns: + tuple: (new_down, new_up, new_alpha) where new_alpha may be None if create_alpha=False + """ + conv2d = len(lora_down.size()) == 4 + lora_down = lora_down.to(process_dtype).to(device) + lora_up = lora_up.to(process_dtype).to(device) + + if conv2d: + out_size, in_rank_up, _, _ = lora_up.shape + in_rank_down, in_size, kh, kw = lora_down.shape + merged = lora_up.reshape(out_size, -1) @ lora_down.reshape(in_rank_down, -1) + else: + merged = lora_up @ lora_down + + U, Vh, actual_rank = decompose_to_rank(merged, new_rank) + + if conv2d: + new_down = Vh.reshape(actual_rank, in_size, kh, kw).cpu() + new_up = U.reshape(out_size, actual_rank, 1, 1).cpu() + else: + new_down = Vh.cpu() + new_up = U.cpu() + + # Only calculate alpha if requested + new_alpha = None + if create_alpha: + orig_rank = lora_down.shape[0] + alpha_val = alpha.item() if isinstance(alpha, torch.Tensor) else (alpha or orig_rank) + new_alpha = float(alpha_val * actual_rank / orig_rank) + + return new_down, new_up, new_alpha + +def apply_dare(weight, drop_rate=0.9): + """ + Applies DARE (Drop and REscale) to a weight matrix. + + Args: + weight: The reconstructed weight matrix (Delta). + drop_rate: Proportion of weights to drop (0.9 = 90% dropped). + + Returns: + The thinned weight matrix with same dtype as input. + """ + if drop_rate <= 0: + return weight + + original_dtype = weight.dtype + weight = weight.float() + + # Create a random mask the same shape as the weight + mask = torch.bernoulli(torch.full_like(weight, 1 - drop_rate)) + + # Apply mask and rescale + rescale_factor = 1 / (1 - drop_rate) + result = (weight * mask) * rescale_factor + + return result.to(original_dtype) + + +def apply_magprune(weight, drop_rate=0.9, epsilon=0.5, row_wise=True): + """ + Applies MagPrune (DELLA-style magnitude-based pruning) to a weight matrix. + + DELLA = Drop and rEscaLe via sampLing with mAgnitude + + Instead of random dropping (DARE), ranks weights by magnitude and assigns + higher drop probability to low-magnitude weights. Each weight gets its own + rescaling factor based on its drop probability. + + Args: + weight: The reconstructed weight matrix (Delta). + drop_rate: Average drop probability (p). Default 0.9 = 90% average drop. + epsilon: Spread of drop probabilities. Higher = more aggressive magnitude bias. + 0.0 = uniform (same as DARE), 1.0 = wide spread. + row_wise: If True, rank within each row (better per DELLA paper). + If False, rank globally across the entire matrix. + + Returns: + The thinned weight matrix with same dtype as input. + + Reference: + "DELLA-Merging: Reducing Interference in Model Merging through Magnitude-Based Sampling" + https://arxiv.org/abs/2406.11617 + """ + if drop_rate <= 0: + return weight + if drop_rate >= 1: + return torch.zeros_like(weight) + + original_dtype = weight.dtype + weight = weight.float() + + # Step 1: Rank weights by magnitude + if row_wise: + # Row-wise ranking (better per DELLA paper) + # For each row, rank elements by their absolute magnitude + n = weight.shape[1] + # argsort twice to get ranks: 0 = highest magnitude, n-1 = lowest + ranks = torch.argsort(torch.argsort(weight.abs(), dim=1, descending=True), dim=1) + # Calculate drop probabilities: p_i = p_min + (epsilon/n) * rank_i + # where p_min = drop_rate - epsilon/2 + p_min = drop_rate - epsilon / 2 + delta = epsilon / n + drop_probs = p_min + delta * ranks + else: + # Global ranking across the entire matrix + flat = weight.flatten() + ranks = torch.argsort(torch.argsort(flat.abs(), descending=True)) + n = len(flat) + p_min = drop_rate - epsilon / 2 + delta = epsilon / n + drop_probs_flat = p_min + delta * ranks + drop_probs = drop_probs_flat.reshape(weight.shape) + + # Clamp probabilities to valid range [0, 1] + drop_probs = drop_probs.clamp(0, 1) + + # Step 2: Sample mask and apply per-weight rescaling + mask = torch.bernoulli(1 - drop_probs) + # Per-weight rescaling factor: 1 / (1 - p_i) + rescale_factor = 1 / (1 - drop_probs) + # Avoid division by zero + rescale_factor = torch.where(drop_probs >= 1, torch.zeros_like(rescale_factor), rescale_factor) + + result = (weight * mask) * rescale_factor + + return result.to(original_dtype) + + +def apply_lambda_scaling(weight, lambda_val=1.0, mode="uniform", min_scale=None, top_k=None): + """ + Applies lambda scaling to a weight matrix with multiple modes. + + DELLA found that constant lambda scaling after merge significantly + improves performance over adaptive scaling. The optimal lambda is + typically found by sweeping values like [0.5, 0.7, 1.0, 1.3, 1.5]. + + Args: + weight: The weight matrix to scale. + lambda_val: Scaling factor. 1.0 = no change. + < 1.0 = dampen, > 1.0 = amplify. + mode: Scaling mode. Options: + - "uniform": Constant scaling (default, original behavior) + - "magnitude": Magnitude-aware scaling - preserves small weights + - "singular": Scale only top-k singular values + - "clipped": Prevents weights from being scaled below threshold + min_scale: For "clipped" mode, minimum scale factor (e.g., 0.5) + top_k: For "singular" mode, number of singular values to scale (None = all) + + Returns: + The scaled weight matrix with same dtype as input. + + Reference: + DELLA-Merging paper, Section 4 "Adaptive vs Constant lambda" + """ + if lambda_val == 1.0: + return weight + + if mode == "uniform": + return weight * lambda_val + + original_dtype = weight.dtype + weight_float = weight.float() + + if mode == "magnitude": + # Magnitude-aware scaling: Scale less aggressively for small weights + # This preserves fine details while still applying scaling to strong weights + with torch.no_grad(): + abs_weight = weight_float.abs() + max_abs = abs_weight.max().clamp(min=1e-8) + + # Normalize to [0, 1], then apply magnitude-aware scaling + # Small weights get scale closer to 1.0, large weights get full lambda + normalized = abs_weight / max_abs + + # Interpolate between 1.0 and lambda based on magnitude + # Small magnitude -> scale closer to 1.0 (preserve) + # Large magnitude -> scale closer to lambda (full effect) + scale_per_weight = 1.0 + (lambda_val - 1.0) * normalized + + result = weight_float * scale_per_weight + return result.to(original_dtype) + + elif mode == "singular": + # Singular value scaling: Scale only the top-k singular values + # This preserves the rank structure while still applying scaling + U, S, Vh = torch.linalg.svd(weight_float, full_matrices=False) + + # Scale singular values + S_scaled = S.clone() + if top_k is None: + S_scaled = S * lambda_val + else: + k = min(top_k, len(S)) + S_scaled[:k] = S[:k] * lambda_val + # Remaining singular values unchanged + + # Reconstruct + result = U @ torch.diag(S_scaled) @ Vh + return result.to(original_dtype) + + elif mode == "clipped": + # Clipped scaling: Ensure small weights aren't scaled below threshold + # min_scale acts as a floor - no weight gets scaled more than this + if min_scale is None: + min_scale = 0.5 + + # For lambda < 1 (reduction), min_scale is the maximum reduction + # For lambda > 1 (amplification), we clamp the lower bound + if lambda_val < 1.0: + # When reducing, don't reduce below min_scale + effective_scale = max(lambda_val, min_scale) + else: + # When amplifying, min_scale prevents small weights from getting too small + # This is more about preserving the minimum magnitude + with torch.no_grad(): + abs_weight = weight_float.abs() + max_abs = abs_weight.max().clamp(min=1e-8) + normalized = abs_weight / max_abs + # Small weights get less amplification + scale_per_weight = 1.0 + (lambda_val - 1.0) * normalized + # Apply minimum scale floor + scale_per_weight = scale_per_weight.clamp(min=min_scale) + result = weight_float * scale_per_weight + return result.to(original_dtype) + + return weight * effective_scale + + else: + raise ValueError(f"Unknown scaling mode: {mode}. Choose from: uniform, magnitude, singular, clipped") + + +def apply_magnitude_scaling(weight, lambda_val=1.0): + """ + Magnitude-aware scaling: Preserves small weights while scaling large ones. + + Small weights are scaled less aggressively than large weights, preventing + fine details from vanishing when lambda < 1.0, and preventing small weights + from becoming too large when lambda > 1.0. + + Args: + weight: The weight matrix to scale. + lambda_val: Scaling factor (1.0 = no change). + + Returns: + The scaled weight matrix with same dtype as input. + """ + return apply_lambda_scaling(weight, lambda_val, mode="magnitude") + + +def apply_singular_value_scaling(weight, lambda_val=1.0, top_k=None): + """ + Singular value scaling: Scales only the top-k singular values. + + This preserves the rank structure of the weight matrix while applying + scaling to the most important directions. The singular values represent + the "energy" in each direction - scaling them preserves the structure. + + Args: + weight: The weight matrix to scale. + lambda_val: Scaling factor (1.0 = no change). + top_k: Number of top singular values to scale (None = all). + + Returns: + The scaled weight matrix with same dtype as input. + """ + return apply_lambda_scaling(weight, lambda_val, mode="singular", top_k=top_k) + + +def apply_clipped_scaling(weight, lambda_val=1.0, min_scale=0.5): + """ + Clipped scaling: Prevents weights from being scaled below a threshold. + + When reducing weights (lambda < 1.0), ensures no weight is scaled below + min_scale. This prevents small weights from vanishing completely. + + Args: + weight: The weight matrix to scale. + lambda_val: Scaling factor (1.0 = no change). + min_scale: Minimum scale factor (default 0.5). + + Returns: + The scaled weight matrix with same dtype as input. + """ + return apply_lambda_scaling(weight, lambda_val, mode="clipped", min_scale=min_scale) + + +# ============================================================================ +# DO-Merge: Magnitude & Direction Decoupling (from AIT/MergeOrthogonalProcess) +# ============================================================================ + +def frobenius_norm_product(U: torch.Tensor, D: torch.Tensor) -> torch.Tensor: + """ + Computes the Frobenius norm of (U @ D) in O(R^2 * d) time instead of O(d^2) time. + norm(U @ D)^2 = tr(D^T U^T U D) = tr(U^T U D D^T) + + This is used in DO-Merge to efficiently compute LoRA magnitudes without + expanding to the full [d_out, d_in] matrix. + + Args: + U: Up projection matrix [d_out, rank] + D: Down projection matrix [rank, d_in] + + Returns: + Frobenius norm of the product U @ D + """ + UT_U = torch.matmul(U.t(), U) + D_DT = torch.matmul(D, D.t()) + return torch.sqrt(torch.abs(torch.trace(torch.matmul(UT_U, D_DT)))) + + +def trace_product(U1: torch.Tensor, D1: torch.Tensor, U2: torch.Tensor, D2: torch.Tensor) -> torch.Tensor: + """ + Computes the trace of (U1 @ D1)^T @ (U2 @ D2) efficiently. + tr(D1^T U1^T U2 D2) = tr(D2 D1^T U1^T U2) + + This measures the alignment between two LoRA weight matrices without + expanding to full size. Used in DO-Merge to compute directional similarity. + + Args: + U1, D1: Up and down matrices of first LoRA + U2, D2: Up and down matrices of second LoRA + + Returns: + Trace value representing alignment (-1 to 1 range approximately) + """ + U1T_U2 = torch.matmul(U1.t(), U2) + D2_D1T = torch.matmul(D2, D1.t()) + return torch.trace(torch.matmul(D2_D1T, U1T_U2)) + + +def decoupled_magnitude_direction_merge( + W_up1: torch.Tensor, W_up2: torch.Tensor, + W_down1: torch.Tensor, W_down2: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Magnitude/Direction Decoupling (DO-Merge style) with Exact Rank Preservation. + + This solves the issue where one LoRA has massive weight magnitudes and overpowers + the other during structure merging, even if their directional vectors are aligned. + + By distributing the calculated DO-Merge scalars directly into the down-projection + matrices, we achieve mathematically perfect concatenation with ZERO data loss, + ZERO approximation error, and instantaneous calculation speed (bypassing SVD). + + Algorithm: + 1. Calculate true magnitudes efficiently using frobenius_norm_product + 2. Calculate direction alignment via trace_product + 3. Compute balanced target magnitude using geometric mean + 4. Calculate scaling factors for each LoRA to achieve exact DO-Merge equation + 5. Apply scalars and concatenate (100% exact math, 0 data loss, 0 SVD) + + Args: + W_up1: Up matrix of first LoRA [d_out, rank1] + W_up2: Up matrix of second LoRA [d_out, rank2] + W_down1: Down matrix of first LoRA [rank1, d_in] + W_down2: Down matrix of second LoRA [rank2, d_in] + + Returns: + tuple: (lora_up_new, lora_down_new, optimal_rank) + - lora_up_new: Concatenated up matrix [d_out, rank1 + rank2] + - lora_down_new: Scaled and concatenated down matrix [rank1 + rank2, d_in] + - optimal_rank: rank1 + rank2 (exact preservation) + + Reference: + DO-Merge: "Decoupling Magnitude and Direction for Efficient Model Merging" + AIT implementation: MergeOrthogonalProcess.py lines 61-100 + """ + # 1. Calculate true magnitudes efficiently without expanding to full matrices + mag1 = frobenius_norm_product(W_up1.float(), W_down1.float()) + 1e-8 + mag2 = frobenius_norm_product(W_up2.float(), W_down2.float()) + 1e-8 + + # 2. Calculate the direction sum norm efficiently + # norm(dir1 + dir2)^2 = 2 + 2 * tr(dir1^T dir2) + cross_trace = trace_product(W_up1.float(), W_down1.float(), W_up2.float(), W_down2.float()) + norm_dir_sum_sq = 2.0 + 2.0 * cross_trace / (mag1 * mag2) + norm_dir_sum_sq = torch.clamp(norm_dir_sum_sq, min=1e-8) + norm_dir_sum = torch.sqrt(norm_dir_sum_sq) + + # 3. Calculate the balanced target magnitude (Geometric Mean) + merged_mag = torch.sqrt(mag1 * mag2) + + # 4. Calculate final scaling factors for each LoRA to achieve the exact DO-Merge equation + global_scalar = merged_mag / norm_dir_sum + + scale1 = global_scalar / mag1 + scale2 = global_scalar / mag2 + + # 5. Apply scalars to the down matrices and concatenate (100% exact math, 0 data loss, 0 SVD) + down1_new = W_down1.float() * scale1 + down2_new = W_down2.float() * scale2 + + lora_up_new = torch.cat([W_up1.float(), W_up2.float()], dim=1) + lora_down_new = torch.cat([down1_new, down2_new], dim=0) + + optimal_rank = lora_up_new.shape[1] # rank1 + rank2 exactly + + return lora_up_new.to(W_up1.dtype), lora_down_new.to(W_down1.dtype), optimal_rank + + +def apply_ties_delta_scaling(delta_weights, densities, lambda_val=1.0): + """ + Applies DELLA-style lambda scaling to TIES-merged deltas. + + After the TIES merge (Trim, Elect, Merge), applies a constant + lambda scaling factor to boost or dampen the merged result. + + Args: + delta_weights: List of reconstructed weight matrices (Deltas) from each LoRA + densities: List of densities (0.0 to 1.0) for trimming + lambda_val: Post-merge scaling factor + + Returns: + Scaled merged weight matrix. + """ + merged = apply_ties(delta_weights, densities) + return apply_lambda_scaling(merged, lambda_val) + +def apply_ties(delta_weights, densities): + """ + Applies TIES (Trim, Elect, and Merge) to multiple weight matrices. + + TIES treats weight deltas like a democratic election to decide which + LoRA gets to control which specific parameter. + + Args: + delta_weights: List of reconstructed weight matrices (Deltas) from each LoRA + densities: List of floats (0.0 to 1.0) for trimming noise per LoRA. + 0.2 = keep top 20% most significant weights, zero out the rest. + + Returns: + A single consensus weight matrix. + + The three stages: + 1. Trim: Keep only top k% magnitudes per LoRA (remove noise) + 2. Elect: Sign consensus - positive vs negative "votes" per parameter + 3. Merge: Disjoint merge - average only values that agree with winning sign + """ + if len(delta_weights) == 0: + raise ValueError("No delta weights provided") + + if len(delta_weights) == 1: + return delta_weights[0] + + # Store original dtype for consistent output + original_dtype = delta_weights[0].dtype + + # ============================== + # Stage 1: Trim - Keep only top k% magnitudes per LoRA + # ============================== + trimmed_deltas = [] + for delta, density in zip(delta_weights, densities): + if density >= 1.0: + trimmed_deltas.append(delta) + continue + + if density <= 0.0: + trimmed_deltas.append(torch.zeros_like(delta)) + continue + + # Convert to float for consistent threshold computation + delta_float = delta.float() + flat_delta = delta_float.abs().flatten() + + # Calculate k (number of elements to keep) + k = int(density * flat_delta.numel()) + if k == 0: + trimmed_deltas.append(torch.zeros_like(delta)) + continue + + # Get threshold value at the k-th position + threshold = torch.topk(flat_delta, k).values[-1] + + # Create mask and apply + mask = (delta_float.abs() >= threshold).to(delta_float.dtype) + trimmed_deltas.append((delta_float * mask).to(delta.dtype)) + + # ============================== + # Stage 2: Elect - Sign Consensus + # ============================== + stacked_deltas = torch.stack(trimmed_deltas) + signs = torch.sign(stacked_deltas) + + # Sum the magnitudes for positive and negative directions + # This creates the "vote count" for each direction + pos_mask = (signs > 0).float() + neg_mask = (signs < 0).float() + + pos_sum = (stacked_deltas.abs() * pos_mask).sum(dim=0) + neg_sum = (stacked_deltas.abs() * neg_mask).sum(dim=0) + + # The winning sign is the one with more total magnitude + final_sign = torch.where(pos_sum >= neg_sum, 1.0, -1.0) + + # ============================== + # Stage 3: Disjoint Merge - Average only agreeing values + # ============================== + # Create mask where LoRA agrees with the winning sign + consensus_mask = (signs == final_sign).float() + + # Mask out disagreeing values and sum across LoRAs + sum_agreed = (stacked_deltas * consensus_mask).sum(dim=0) + + # Count how many LoRAs contributed to each parameter's agreement + count_agreed = consensus_mask.sum(dim=0).clamp(min=1.0) + + # Average only the values that agreed with the winning sign + result = sum_agreed / count_agreed + + return result.to(original_dtype) diff --git a/utility/ltx/utils/__init__.py b/utility/ltx/utils/__init__.py new file mode 100644 index 0000000..6ddc486 --- /dev/null +++ b/utility/ltx/utils/__init__.py @@ -0,0 +1,22 @@ +"""LTX utilities for I/O and naming.""" + +from .io import load_lora, save_lora +from .naming import ( + LORA_DOWN_UP_FORMATS, + is_audio_layer, + is_video_ff_layer, + is_video_only_layer, + normalize_lora_key, + build_block_map +) + +__all__ = [ + 'load_lora', + 'save_lora', + 'LORA_DOWN_UP_FORMATS', + 'is_audio_layer', + 'is_video_ff_layer', + 'is_video_only_layer', + 'normalize_lora_key', + 'build_block_map' +] diff --git a/utility/ltx/utils/io.py b/utility/ltx/utils/io.py new file mode 100644 index 0000000..c843e8b --- /dev/null +++ b/utility/ltx/utils/io.py @@ -0,0 +1,65 @@ +import torch +from safetensors.torch import save_file, safe_open +from pathlib import Path +from typing import Optional, Dict + + +def load_lora(path, device="cpu"): + """Load LoRA from safetensors or pytorch file.""" + path = Path(path) + if path.suffix == ".safetensors": + with safe_open(path, framework="pt", device=device) as f: + sd = {k: f.get_tensor(k) for k in f.keys()} + metadata = f.metadata() + else: + sd = torch.load(path, map_location=device) + if isinstance(sd, dict) and "state_dict" in sd: + sd = sd["state_dict"] + metadata = None + return sd, metadata + + +def save_lora(sd, output_path, metadata=None, dtype=None, output_format="lora"): + """ + Save state dict to safetensors with optional conversion. + + Args: + sd: State dict to save + output_path: Output file path + metadata: Optional metadata dict + dtype: Optional dtype for conversion + output_format: Format to save as - "lora" or "lokr" + """ + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Convert to LoKR format if requested + if output_format == "lokr": + from ltx_merger.utils.lokr import convert_lora_to_lokr + sd = convert_lora_to_lokr(sd) + + # Update metadata for LoKR + from ltx_merger.utils.lokr import get_lokr_metadata + factor = None + conv_dim = None + + # Try to extract factor from metadata or state dict + if metadata: + if "ss_network_dim" in metadata: + factor = int(metadata.get("ss_network_dim", factor)) + if "lokr_conv_dim" in metadata: + conv_dim = int(metadata.get("lokr_conv_dim", conv_dim)) + + metadata = get_lokr_metadata(factor=factor, conv_dim=conv_dim, original_metadata=metadata) + + processed_sd = {} + for k, v in sd.items(): + if isinstance(v, torch.Tensor): + v = v.cpu() + if dtype is not None and v.dtype.is_floating_point: + v = v.to(dtype) + processed_sd[k] = v + else: + processed_sd[k] = v + + save_file(processed_sd, str(output_path), metadata=metadata) \ No newline at end of file diff --git a/utility/ltx/utils/naming.py b/utility/ltx/utils/naming.py new file mode 100644 index 0000000..a754411 --- /dev/null +++ b/utility/ltx/utils/naming.py @@ -0,0 +1,203 @@ +import torch + +LORA_DOWN_UP_FORMATS = [ + ("lora_down", "lora_up"), # sd-scripts LoRA + ("lora_A", "lora_B"), # PEFT LoRA + ("down", "up"), # ControlLoRA +] + +def is_audio_layer(key): + """Check if a key corresponds to an audio/video layer.""" + # Use dot-notation checks to avoid false positives (e.g., "audio" in other contexts) + key_lower = key.lower() + return ( + ".audio_attn" in key_lower or + "audio_ff" in key_lower or + "audio_to_video_attn" in key_lower or + "video_to_audio_attn" in key_lower + ) + + +def is_audio_ff_layer(key: str) -> bool: + """Detect audio FF layers specifically (audio_ff in key).""" + return ".audio_ff." in key + + +def is_video_ff_layer(key: str) -> bool: + """Detect video FF layers (ff layers that are NOT audio_ff). + + Note: to_gate_logits is part of attention layers (not FF), so we + explicitly exclude it from FF classification. + """ + # Check for actual FF/MLP layer patterns (using dot notation to avoid false matches) + key_lower = key.lower() + # Exclude to_gate_logits (part of attention, not FF) + if "to_gate_logits" in key_lower: + return False + # Look for actual FF/MLP patterns (e.g., ".ff.", ".feed_forward.", ".mlp.", ".gate.") + # Using dot notation avoids false matches like "ff" in "diffusion_model" + # Also exclude audio_ff since those are audio layers, not video FF + has_ff_pattern = (".ff." in key_lower or ".feed_forward." in key_lower or ".mlp." in key_lower or ".gate." in key_lower) + is_not_audio_ff = ".audio_ff." not in key_lower + return has_ff_pattern and is_not_audio_ff + + +# ============================================================================ +# Semantic Layer Detection (for AIT-style smart merging) +# ============================================================================ + +def is_cross_attention_layer(key: str) -> bool: + """ + Detect cross-attention layers (prompt routing). + + Cross-attention layers handle text-to-image/video conditioning and act + as "trigger routers" - they control which prompts activate which features. + + These benefit from Bilateral Subspace Orthogonalization (BSO) to prevent + trigger conflicts between merged LoRAs. + """ + key_lower = key.lower() + return any( + indicator in key_lower + for indicator in ["attn2", "txt_attn", "context", "txt_in", "enc"] + ) + + +def is_temporal_layer(key: str) -> bool: + """ + Detect temporal/motion layers. + + Temporal layers handle video motion and timing. Like audio, they act as + style controllers and benefit from BSO to prevent motion style conflicts. + """ + key_lower = key.lower() + return any( + indicator in key_lower + for indicator in ["temp", "time", "video", "motion"] + ) + + +def is_mlp_layer(key: str) -> bool: + """ + Detect MLP/structural layers. + + MLP layers handle the actual feature transformation (not routing). They + benefit from Magnitude/Direction Decoupling (DO-Merge) to prevent one + LoRA's magnitude from overpowering another. + + Note: to_gate_logits is part of attention layers (not MLP/FF), so we + explicitly exclude it from MLP classification. + """ + key_lower = key.lower() + # Exclude to_gate_logits (part of attention, not MLP/FF) + if "to_gate_logits" in key_lower: + return False + # Look for actual FF/MLP patterns (using dot notation to avoid false matches) + # Note: audio_ff is also an MLP/FF layer, so we include it + return any( + f".{indicator}." in key_lower + for indicator in ["ff", "audio_ff", "feed_forward", "mlp", "gate"] + ) + + +def is_video_only_layer(key: str) -> bool: + """ + Detect video-only layers (non-FF, non-audio). + + Video-only layers are cross-attention, temporal, and other structural + layers that are NOT audio layers and NOT FF/MLP layers. + """ + return not is_audio_layer(key) and not is_mlp_layer(key) + + +def get_layer_semantic_type(key: str) -> str: + """ + Determine the semantic type of a layer for smart merging. + + Returns: + "cross_attn": Cross-attention (prompt routing) → use BSO + "audio": Audio layers → use BSO + "temporal": Temporal/motion layers → use BSO + "mlp": Structural/MLP layers → use DO-Merge + "other": Other layers → default behavior + """ + if is_audio_layer(key): + return "audio" + if is_cross_attention_layer(key): + return "cross_attn" + if is_temporal_layer(key): + return "temporal" + if is_mlp_layer(key): + return "mlp" + return "other" + +def normalize_lora_key(key): + """Normalize a LoRA key to a standard format for matching.""" + normalized = key + for fmt in LORA_DOWN_UP_FORMATS: + normalized = normalized.replace(f".{fmt[0]}", ".lora_down") + normalized = normalized.replace(f".{fmt[1]}", ".lora_up") + return normalized + +def get_layer_info(lora_sd): + """Analyze LoRA layers and count different types.""" + layer_counts = { + "attn1": 0, "attn2": 0, "ff": 0, "audio": 0, "other": 0, "total": 0 + } + blocks = set() + + for key in lora_sd.keys(): + if key.endswith(".alpha"): continue + + key_parts = key.split(".") + for fmt_down, _ in LORA_DOWN_UP_FORMATS: + if len(key_parts) >= 2 and fmt_down == key_parts[-2]: + block_name = ".".join(key_parts[:-2]) + if block_name not in blocks: + blocks.add(block_name) + if is_audio_layer(block_name): + layer_counts["audio"] += 1 + elif "attn1" in block_name: + layer_counts["attn1"] += 1 + elif "attn2" in block_name: + layer_counts["attn2"] += 1 + elif any(x in block_name for x in ["ff", "feed_forward"]): + layer_counts["ff"] += 1 + else: + layer_counts["other"] += 1 + break + + layer_counts["total"] = len(blocks) + return layer_counts + +def build_block_map(lora_sd): + """ + Build a block map from a single LoRA state dict. + Returns dict mapping block_name -> block_data with down/up weights. + Used by both pruner.py and merger.py. + """ + blocks = {} + for key in lora_sd.keys(): + if key.endswith(".alpha"): continue + + key_parts = key.split(".") + for fmt_down, fmt_up in LORA_DOWN_UP_FORMATS: + if len(key_parts) >= 2 and fmt_down == key_parts[-2]: + block_name = ".".join(key_parts[:-2]) + + # Determine suffix (.weight or empty) + weight_suffix = "." + key_parts[-1] if key_parts[-1] in ["weight", "bias"] else "" + up_key = f"{block_name}.{fmt_up}{weight_suffix}" + + if up_key in lora_sd: + blocks[block_name] = { + "down": lora_sd[key], + "up": lora_sd[up_key], + "alpha": lora_sd.get(f"{block_name}.alpha"), + "down_name": key, + "up_name": up_key, + "is_conv2d": len(lora_sd[key].shape) == 4, + "is_audio": is_audio_layer(block_name) + } + break + return blocks \ No newline at end of file diff --git a/utility/ltx/utils/progress.py b/utility/ltx/utils/progress.py new file mode 100644 index 0000000..ee5cbdc --- /dev/null +++ b/utility/ltx/utils/progress.py @@ -0,0 +1,81 @@ +"""Unified progress bar for all LTX merger operations.""" + +import sys +import time + +class ProgressBar: + """Simple, reusable progress bar for console output.""" + + def __init__(self, total, desc="Processing", width=40): + self.total = total + self.current = 0 + self.desc = desc + self.width = width + self.start_time = None + + def start(self): + """Initialize the progress bar.""" + self.start_time = time.time() + self.update(0) + + def update(self, n=1): + """Advance progress by n steps.""" + self.current += n + self._draw() + + def set(self, value): + """Set progress to specific value.""" + self.current = value + self._draw() + + def _draw(self): + """Draw the progress bar.""" + if self.total == 0: + return + + percent = min(1.0, self.current / self.total) + filled = int(self.width * percent) + bar = "█" * filled + "░" * (self.width - filled) + + # Calculate elapsed time + elapsed = time.time() - self.start_time if self.start_time else 0 + if self.current > 0 and percent < 1.0: + eta = elapsed / percent * (1 - percent) + time_str = f" {int(eta)}s" + else: + time_str = "" + + # Clear line and draw + sys.stdout.write(f"\r{self.desc}: [{bar}] {self.current}/{self.total} ({percent*100:.1f}%){time_str}") + sys.stdout.flush() + + # New line when complete + if self.current >= self.total: + sys.stdout.write(f"\n{self.desc} completed in {int(elapsed)}s\n") + sys.stdout.flush() + + def finish(self): + """Force completion of the progress bar.""" + self.current = self.total + self._draw() + + +def progress_iterator(items, desc="Processing"): + """Wrap an iterable with a progress bar. + + Usage: + for item in progress_iterator(my_list, "Loading"): + process(item) + """ + total = len(items) if hasattr(items, '__len__') else None + if total is None: + # No progress bar for unknown length iterables + for item in items: + yield item + return + + bar = ProgressBar(total, desc) + bar.start() + for item in items: + yield item + bar.update(1)