AB overlap
This commit is contained in:
JJ
@@ -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 *
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
+1102
-84
File diff suppressed because it is too large
Load Diff
@@ -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'
|
||||
]
|
||||
@@ -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'
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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'
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user