573 lines
26 KiB
Python
573 lines
26 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import folder_paths
|
|
import comfy.utils
|
|
import comfy.sd
|
|
from safetensors.torch import load_file
|
|
import re
|
|
import os
|
|
|
|
# --- Helper Functions ---
|
|
|
|
def convert_alpha_key_to_basis_base(key):
|
|
"""Alpha key (te1.text.model...) -> Basis base name"""
|
|
return key.replace("_", ".")
|
|
|
|
def convert_basis_base_to_comfy_key(basis_base_key):
|
|
"""Basis base name (te1.text.model...) -> ComfyUI LoRA key (lora_te1_text_model...)"""
|
|
if basis_base_key.startswith("unet."):
|
|
new_key = "lora_unet_" + basis_base_key[5:]
|
|
elif basis_base_key.startswith("te1."):
|
|
new_key = "lora_te1_" + basis_base_key[4:]
|
|
elif basis_base_key.startswith("te2."):
|
|
new_key = "lora_te2_" + basis_base_key[4:]
|
|
elif basis_base_key.startswith("te."):
|
|
new_key = "lora_te_" + basis_base_key[3:]
|
|
else:
|
|
new_key = "lora_unet_" + basis_base_key # Fallback
|
|
return new_key.replace(".", "_")
|
|
|
|
def map_model_keys_to_lora_keys(model, clip):
|
|
"""Scan all layers in the model and create a map to LoRA keys"""
|
|
shape_map = {}
|
|
|
|
# --- 1. UNet Mapping ---
|
|
unet_root = model.model.diffusion_model
|
|
for name, module in unet_root.named_modules():
|
|
if module.__class__.__name__ in ["Linear", "Conv2d", "LoRACompatibleLinear", "LoRACompatibleConv"]:
|
|
lora_key_base = "lora_unet_" + name.replace(".", "_")
|
|
if hasattr(module, "weight"):
|
|
shape_map[lora_key_base] = module.weight.shape
|
|
|
|
# --- 2. CLIP Mapping ---
|
|
def scan_te(te_model, prefix):
|
|
for name, module in te_model.named_modules():
|
|
if module.__class__.__name__ in ["Linear", "Conv2d", "LoRACompatibleLinear", "LoRACompatibleConv"]:
|
|
if hasattr(module, "weight"):
|
|
clean_name = name
|
|
if clean_name.startswith("transformer."):
|
|
clean_name = clean_name.replace("transformer.", "")
|
|
lora_key_base = prefix + clean_name.replace(".", "_")
|
|
shape_map[lora_key_base] = module.weight.shape
|
|
|
|
raw_key_base = prefix + name.replace(".", "_")
|
|
if raw_key_base not in shape_map:
|
|
shape_map[raw_key_base] = module.weight.shape
|
|
|
|
if hasattr(clip.cond_stage_model, "clip_l"):
|
|
scan_te(clip.cond_stage_model.clip_l, "lora_te1_")
|
|
if hasattr(clip.cond_stage_model, "clip_g"):
|
|
scan_te(clip.cond_stage_model.clip_g, "lora_te2_")
|
|
if hasattr(clip.cond_stage_model, "transformer"):
|
|
scan_te(clip.cond_stage_model.transformer, "lora_te_")
|
|
|
|
return shape_map
|
|
|
|
def resize_lora_weight(weight, expected_param_count, is_up=False):
|
|
"""
|
|
Adjusts the LoRA weight rank to match the Subspace Basis using Slicing.
|
|
"""
|
|
current_count = weight.numel()
|
|
if current_count == expected_param_count:
|
|
return weight
|
|
|
|
shape = weight.shape
|
|
|
|
# Identify Rank Dimension and Target Rank
|
|
if len(shape) == 2: # Linear
|
|
if not is_up: # Down: (Rank, In)
|
|
rank_dim = 0
|
|
other_dim = shape[1]
|
|
current_rank = shape[0]
|
|
required_rank = expected_param_count // other_dim
|
|
else: # Up: (Out, Rank)
|
|
rank_dim = 1
|
|
other_dim = shape[0]
|
|
current_rank = shape[1]
|
|
required_rank = expected_param_count // other_dim
|
|
elif len(shape) == 4: # Conv
|
|
if not is_up: # Down: (Rank, In, K, K)
|
|
rank_dim = 0
|
|
other_dims = shape[1] * shape[2] * shape[3]
|
|
current_rank = shape[0]
|
|
required_rank = expected_param_count // other_dims
|
|
else: # Up: (Out, Rank, 1, 1)
|
|
rank_dim = 1
|
|
other_dims = shape[0] * shape[2] * shape[3]
|
|
current_rank = shape[1]
|
|
required_rank = expected_param_count // other_dims
|
|
else:
|
|
return weight # Unknown shape
|
|
|
|
if required_rank == current_rank:
|
|
return weight
|
|
|
|
# Case 1: Rank Insufficient (Padding)
|
|
if required_rank > current_rank:
|
|
pad_size = required_rank - current_rank
|
|
pad_shape = list(shape)
|
|
pad_shape[rank_dim] = pad_size
|
|
padding = torch.zeros(pad_shape, dtype=weight.dtype, device=weight.device)
|
|
|
|
if rank_dim == 0:
|
|
new_weight = torch.cat([weight, padding], dim=0)
|
|
else:
|
|
new_weight = torch.cat([weight, padding], dim=1)
|
|
|
|
# Case 2: Rank Excess (Slicing)
|
|
else:
|
|
if rank_dim == 0:
|
|
new_weight = weight[:required_rank, ...]
|
|
else:
|
|
# Assuming rank is dim 1
|
|
if len(shape) == 2:
|
|
new_weight = weight[:, :required_rank]
|
|
elif len(shape) == 4:
|
|
new_weight = weight[:, :required_rank, :, :]
|
|
else:
|
|
new_weight = weight
|
|
|
|
return new_weight
|
|
|
|
# --- Class Definitions ---
|
|
|
|
class UniversalLoRALoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"universal_lora_name": (folder_paths.get_filename_list("loras"), )}}
|
|
|
|
RETURN_TYPES = ("UNIVERSAL_SUBSPACE",)
|
|
FUNCTION = "load_subspace"
|
|
CATEGORY = "UniversalSubspace"
|
|
|
|
def load_subspace(self, universal_lora_name):
|
|
lora_path = folder_paths.get_full_path("loras", universal_lora_name)
|
|
state_dict = load_file(lora_path)
|
|
print(f"DEBUG: Loaded Basis/LoRA. Count: {len(state_dict)}")
|
|
return (state_dict,)
|
|
|
|
class LoadUniversalWeights:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"weight_file": (folder_paths.get_filename_list("loras"), )}}
|
|
|
|
RETURN_TYPES = ("UNIVERSAL_WEIGHTS",)
|
|
FUNCTION = "load_weights"
|
|
CATEGORY = "UniversalSubspace"
|
|
|
|
def load_weights(self, weight_file):
|
|
path = folder_paths.get_full_path("loras", weight_file)
|
|
weights = load_file(path)
|
|
print(f"DEBUG: Loaded Weights. Count: {len(weights)}")
|
|
return (weights,)
|
|
|
|
class ApplyUniversalWeights:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"subspace": ("UNIVERSAL_SUBSPACE",),
|
|
"universal_weights": ("UNIVERSAL_WEIGHTS",),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "apply_weights"
|
|
CATEGORY = "UniversalSubspace"
|
|
|
|
def apply_weights(self, model, clip, subspace, universal_weights, strength_model, strength_clip):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print("DEBUG: Mapping model keys for Universal Weights...")
|
|
shape_map = map_model_keys_to_lora_keys(model, clip)
|
|
|
|
reconstructed_lora = {}
|
|
count = 0
|
|
|
|
for key, alpha in universal_weights.items():
|
|
if not key.endswith(".alpha"): continue
|
|
|
|
alpha_base = key.replace(".alpha", "")
|
|
basis_base = convert_alpha_key_to_basis_base(alpha_base)
|
|
down_prefix = basis_base + ".lora_down.weight"
|
|
up_prefix = basis_base + ".lora_up.weight"
|
|
|
|
if (down_prefix + ".mean") not in subspace or (up_prefix + ".mean") not in subspace:
|
|
continue
|
|
|
|
try:
|
|
a = alpha.to(device, dtype=torch.float32)
|
|
|
|
# --- Down Reconstruction ---
|
|
mu_d = subspace[down_prefix + ".mean"].to(device, dtype=torch.float32)
|
|
V_d = subspace[down_prefix + ".basis"].to(device, dtype=torch.float32)
|
|
w_down_flat = mu_d + torch.matmul(V_d, a)
|
|
|
|
# --- Up Reconstruction ---
|
|
mu_u = subspace[up_prefix + ".mean"].to(device, dtype=torch.float32)
|
|
V_u = subspace[up_prefix + ".basis"].to(device, dtype=torch.float32)
|
|
w_up_flat = mu_u + torch.matmul(V_u, a)
|
|
|
|
# --- Shape Matching ---
|
|
comfy_key_base = convert_basis_base_to_comfy_key(basis_base)
|
|
target_shape = shape_map.get(comfy_key_base, None)
|
|
final_key_name = comfy_key_base
|
|
|
|
if target_shape is None:
|
|
if "_text_model_" in comfy_key_base:
|
|
alt_key = comfy_key_base.replace("_text_model_", "_")
|
|
target_shape = shape_map.get(alt_key, None)
|
|
if target_shape is not None: final_key_name = alt_key
|
|
|
|
if target_shape is None: continue
|
|
|
|
if len(target_shape) == 2:
|
|
out_dim, in_dim = target_shape
|
|
if w_down_flat.numel() % in_dim != 0: continue
|
|
rank = w_down_flat.numel() // in_dim
|
|
w_down = w_down_flat.reshape(rank, in_dim)
|
|
w_up = w_up_flat.reshape(out_dim, rank)
|
|
elif len(target_shape) == 4:
|
|
out_dim, in_dim, kh, kw = target_shape
|
|
ks = kh * kw
|
|
if w_down_flat.numel() % (in_dim * ks) != 0: continue
|
|
rank = w_down_flat.numel() // (in_dim * ks)
|
|
w_down = w_down_flat.reshape(rank, in_dim, kh, kw)
|
|
w_up = w_up_flat.reshape(out_dim, rank, 1, 1)
|
|
else: continue
|
|
|
|
reconstructed_lora[final_key_name + ".lora_down.weight"] = w_down.cpu()
|
|
reconstructed_lora[final_key_name + ".lora_up.weight"] = w_up.cpu()
|
|
count += 1
|
|
|
|
except Exception as e:
|
|
print(f"Error reconstructing {basis_base}: {e}")
|
|
|
|
print(f"DEBUG: Universal LoRA Applied Layers: {count}")
|
|
if count == 0: return (model, clip)
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, reconstructed_lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|
|
|
|
class ApplyUniversalLoRA:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"subspace": ("UNIVERSAL_SUBSPACE",),
|
|
"target_lora_name": (folder_paths.get_filename_list("loras"), ),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "apply_universal"
|
|
CATEGORY = "UniversalSubspace"
|
|
|
|
def apply_universal(self, model, clip, subspace, target_lora_name, strength_model, strength_clip):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# 1. Load Target LoRA
|
|
target_path = folder_paths.get_full_path("loras", target_lora_name)
|
|
target_lora = load_file(target_path)
|
|
print(f"DEBUG: Loaded Target LoRA: {target_lora_name} ({len(target_lora)} keys)")
|
|
|
|
final_lora = {}
|
|
|
|
# 2. Identify Basis Keys
|
|
basis_keys = set()
|
|
for key in subspace.keys():
|
|
if key.endswith(".lora_down.weight.basis"):
|
|
base = key.replace(".lora_down.weight.basis", "")
|
|
basis_keys.add(base)
|
|
|
|
count_approx = 0
|
|
|
|
# 3. Process each basis block
|
|
for basis_base in basis_keys:
|
|
comfy_key_base = convert_basis_base_to_comfy_key(basis_base)
|
|
down_key = comfy_key_base + ".lora_down.weight"
|
|
up_key = comfy_key_base + ".lora_up.weight"
|
|
|
|
# Fallback search for key names
|
|
if down_key not in target_lora:
|
|
if "_text_model_" in comfy_key_base:
|
|
alt = comfy_key_base.replace("_text_model_", "_")
|
|
if (alt + ".lora_down.weight") in target_lora:
|
|
down_key = alt + ".lora_down.weight"
|
|
up_key = alt + ".lora_up.weight"
|
|
|
|
if down_key in target_lora and up_key in target_lora:
|
|
try:
|
|
down_prefix = basis_base + ".lora_down.weight"
|
|
up_prefix = basis_base + ".lora_up.weight"
|
|
|
|
if (down_prefix + ".mean") in subspace and (up_prefix + ".mean") in subspace:
|
|
# --- Prepare Data (Joint Projection) ---
|
|
mu_d = subspace[down_prefix + ".mean"].to(device, dtype=torch.float32)
|
|
V_d = subspace[down_prefix + ".basis"].to(device, dtype=torch.float32)
|
|
mu_u = subspace[up_prefix + ".mean"].to(device, dtype=torch.float32)
|
|
V_u = subspace[up_prefix + ".basis"].to(device, dtype=torch.float32)
|
|
|
|
w_down_raw = target_lora[down_key].to(device, dtype=torch.float32)
|
|
w_up_raw = target_lora[up_key].to(device, dtype=torch.float32)
|
|
|
|
w_down = resize_lora_weight(w_down_raw, mu_d.shape[0], is_up=False)
|
|
w_up = resize_lora_weight(w_up_raw, mu_u.shape[0], is_up=True)
|
|
|
|
# Joint System for Least Squares
|
|
target_vec_d = w_down.flatten() - mu_d
|
|
target_vec_u = w_up.flatten() - mu_u
|
|
|
|
joint_target = torch.cat([target_vec_d, target_vec_u], dim=0)
|
|
joint_basis = torch.cat([V_d, V_u], dim=0)
|
|
|
|
result = torch.linalg.lstsq(joint_basis, joint_target)
|
|
alpha = result.solution
|
|
|
|
w_down_recon_flat = mu_d + torch.matmul(V_d, alpha)
|
|
w_up_recon_flat = mu_u + torch.matmul(V_u, alpha)
|
|
|
|
w_down_recon = w_down_recon_flat.reshape(w_down.shape)
|
|
w_up_recon = w_up_recon_flat.reshape(w_up.shape)
|
|
|
|
final_lora[down_key] = w_down_recon.cpu()
|
|
final_lora[up_key] = w_up_recon.cpu()
|
|
|
|
count_approx += 1
|
|
except Exception as e:
|
|
print(f"Warning: Failed to process Joint {basis_base}: {e}")
|
|
|
|
print(f"DEBUG: Joint Subspace Application - Total Reconstructed Pairs: {count_approx}")
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, final_lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|
|
|
|
class ApplyNullLaLoRA:
|
|
"""
|
|
Node to apply Null-LoRA / LaLoRA weights trained with network_null_lalora_bata.py.
|
|
Fully supports:
|
|
- Null Space (Frozen) + Universal Basis (Trainable)
|
|
- Standard Scaling (s)
|
|
- DoRA Magnitude Scaling (m)
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"subspace": ("UNIVERSAL_SUBSPACE",),
|
|
"null_lalora_weights": ("UNIVERSAL_WEIGHTS",),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "apply_null_lalora"
|
|
CATEGORY = "UniversalSubspace"
|
|
|
|
def apply_null_lalora(self, model, clip, subspace, null_lalora_weights, strength_model, strength_clip):
|
|
device = torch.device("cpu") # Process weights on CPU to save VRAM
|
|
print("DEBUG: Mapping model keys for Null-LaLoRA (DoRA/Standard)...")
|
|
shape_map = map_model_keys_to_lora_keys(model, clip)
|
|
|
|
# --- Pre-calculate Subspace Key Map ---
|
|
subspace_lookup = {}
|
|
for k in subspace.keys():
|
|
if k.endswith(".lora_down.weight.mean"):
|
|
basis_base = k.replace(".lora_down.weight.mean", "")
|
|
safe_key = basis_base.replace(".", "_")
|
|
subspace_lookup[safe_key] = basis_base
|
|
|
|
reconstructed_lora = {}
|
|
count = 0
|
|
|
|
# Identify layer base names from the loaded weights
|
|
layer_bases = set()
|
|
for key in null_lalora_weights.keys():
|
|
if key.endswith(".alpha_down"):
|
|
layer_bases.add(key.replace(".alpha_down", ""))
|
|
elif key.endswith(".alpha") and not key.endswith(".alpha_up"):
|
|
layer_bases.add(key.replace(".alpha", ""))
|
|
|
|
for base_name in layer_bases:
|
|
# 1. Resolve Keys
|
|
matched_subspace_base = subspace_lookup.get(base_name, None)
|
|
if matched_subspace_base is None:
|
|
continue
|
|
|
|
comfy_key_base = convert_basis_base_to_comfy_key(matched_subspace_base)
|
|
|
|
# Check target shape
|
|
target_shape = shape_map.get(comfy_key_base, None)
|
|
if target_shape is None:
|
|
if "_text_model_" in comfy_key_base:
|
|
alt_key = comfy_key_base.replace("_text_model_", "_")
|
|
if alt_key in shape_map:
|
|
comfy_key_base = alt_key
|
|
target_shape = shape_map[alt_key]
|
|
|
|
if target_shape is None:
|
|
continue
|
|
|
|
try:
|
|
# --- Retrieve Basis ---
|
|
mu_d = subspace[matched_subspace_base + ".lora_down.weight.mean"].to(device, dtype=torch.float32)
|
|
V_d = subspace[matched_subspace_base + ".lora_down.weight.basis"].to(device, dtype=torch.float32)
|
|
mu_u = subspace[matched_subspace_base + ".lora_up.weight.mean"].to(device, dtype=torch.float32)
|
|
V_u = subspace[matched_subspace_base + ".lora_up.weight.basis"].to(device, dtype=torch.float32)
|
|
|
|
# --- Retrieve Alpha & Reconstruct Trainable ---
|
|
w_down_train_flat = None
|
|
w_up_train_flat = None
|
|
half_rank = 0
|
|
|
|
if (base_name + ".alpha_down") in null_lalora_weights:
|
|
alpha_d = null_lalora_weights[base_name + ".alpha_down"].to(device, dtype=torch.float32)
|
|
alpha_u = null_lalora_weights[base_name + ".alpha_up"].to(device, dtype=torch.float32)
|
|
|
|
half_rank = alpha_d.shape[0]
|
|
# Resize basis if rank mismatches
|
|
if V_d.shape[1] > half_rank: V_d = V_d[:, :half_rank]
|
|
if V_u.shape[1] > half_rank: V_u = V_u[:, :half_rank]
|
|
|
|
# Transpose logic matches network_null_lalora_bata.py
|
|
# w_down = (alpha_d @ basis_d.T) + mean
|
|
w_down_train_flat = (alpha_d @ V_d.t()) + mu_d.view(1, -1)
|
|
# w_up = (basis_u @ alpha_u) + mean
|
|
w_up_train_flat = (V_u @ alpha_u) + mu_u.view(-1, 1)
|
|
|
|
elif (base_name + ".alpha") in null_lalora_weights:
|
|
# Legacy scalar alpha support
|
|
alpha = null_lalora_weights[base_name + ".alpha"].to(device, dtype=torch.float32)
|
|
half_rank = alpha.shape[0]
|
|
if V_d.shape[1] > half_rank: V_d = V_d[:, :half_rank]
|
|
if V_u.shape[1] > half_rank: V_u = V_u[:, :half_rank]
|
|
|
|
w_down_train_flat = mu_d + torch.matmul(V_d, alpha)
|
|
w_up_train_flat = mu_u + torch.matmul(V_u, alpha)
|
|
else:
|
|
continue
|
|
|
|
# --- Retrieve Frozen Null Parts ---
|
|
# Priority: 1. Weights File (if saved), 2. Subspace File
|
|
frozen_down = None
|
|
frozen_up = None
|
|
|
|
# Note on naming:
|
|
# In training script: null_down=(Out, Rank), null_up=(Rank, In)
|
|
# In forward pass: frozen_down=null_up(Rank, In), frozen_up=null_down(Out, Rank)
|
|
if (base_name + ".null_up") in null_lalora_weights:
|
|
frozen_down = null_lalora_weights[base_name + ".null_up"].to(device, dtype=torch.float32)
|
|
frozen_up = null_lalora_weights[base_name + ".null_down"].to(device, dtype=torch.float32)
|
|
elif (matched_subspace_base + ".lora_up.weight.null") in subspace:
|
|
frozen_down = subspace[matched_subspace_base + ".lora_up.weight.null"].to(device, dtype=torch.float32)
|
|
frozen_up = subspace[matched_subspace_base + ".lora_down.weight.null"].to(device, dtype=torch.float32)
|
|
|
|
if frozen_down is None or frozen_up is None:
|
|
continue
|
|
|
|
# Ensure Frozen rank matches half_rank (Slice if necessary)
|
|
if frozen_down.shape[0] > half_rank:
|
|
frozen_down = frozen_down[:half_rank, ...]
|
|
|
|
if frozen_up.shape[1] > half_rank and len(frozen_up.shape) == 2:
|
|
frozen_up = frozen_up[:, :half_rank]
|
|
elif len(frozen_up.shape) == 4 and frozen_up.shape[1] > half_rank:
|
|
frozen_up = frozen_up[:, :half_rank, :, :]
|
|
|
|
# --- Determine Mode: Standard (s) or DoRA (m) ---
|
|
use_dora = False
|
|
if (base_name + ".m") in null_lalora_weights:
|
|
use_dora = True
|
|
scale_vec = null_lalora_weights[base_name + ".m"].to(device, dtype=torch.float32)
|
|
elif (base_name + ".s") in null_lalora_weights:
|
|
use_dora = False
|
|
scale_vec = null_lalora_weights[base_name + ".s"].to(device, dtype=torch.float32)
|
|
s1 = scale_vec[:half_rank]
|
|
s2 = scale_vec[half_rank:]
|
|
else:
|
|
# Fallback default
|
|
scale_vec = torch.ones(target_shape[0], device=device, dtype=torch.float32)
|
|
use_dora = True # Treat as identity DoRA
|
|
|
|
# --- Reshape & Combine ---
|
|
if len(target_shape) == 2:
|
|
# Linear: (Out, In)
|
|
out_dim, in_dim = target_shape
|
|
w_down_train = w_down_train_flat.reshape(half_rank, in_dim)
|
|
w_up_train = w_up_train_flat.reshape(out_dim, half_rank)
|
|
|
|
f_down = frozen_down.reshape(half_rank, in_dim)
|
|
f_up = frozen_up.reshape(out_dim, half_rank)
|
|
|
|
if use_dora:
|
|
# DoRA Logic: Concatenate then Apply Magnitude
|
|
final_down = torch.cat([f_down, w_down_train], dim=0) # (Rank*2, In)
|
|
final_up = torch.cat([w_up_train, f_up], dim=1) # (Out, Rank*2)
|
|
final_up = final_up * scale_vec.view(-1, 1)
|
|
else:
|
|
# Standard Logic: Apply scales at bottleneck
|
|
w_down_scaled = w_down_train * s2.view(-1, 1)
|
|
final_down = torch.cat([f_down, w_down_scaled], dim=0)
|
|
|
|
w_up_scaled = w_up_train * s1.view(1, -1)
|
|
final_up = torch.cat([w_up_scaled, f_up], dim=1)
|
|
|
|
elif len(target_shape) == 4:
|
|
# Conv2d: (Out, In, kh, kw)
|
|
out_dim, in_dim, kh, kw = target_shape
|
|
w_down_train = w_down_train_flat.reshape(half_rank, in_dim, kh, kw)
|
|
w_up_train = w_up_train_flat.reshape(out_dim, half_rank, 1, 1)
|
|
|
|
f_down = frozen_down.reshape(half_rank, in_dim, kh, kw)
|
|
f_up = frozen_up.reshape(out_dim, half_rank, 1, 1)
|
|
|
|
if use_dora:
|
|
final_down = torch.cat([f_down, w_down_train], dim=0)
|
|
final_up = torch.cat([w_up_train, f_up], dim=1)
|
|
final_up = final_up * scale_vec.view(-1, 1, 1, 1)
|
|
else:
|
|
w_down_scaled = w_down_train * s2.view(-1, 1, 1, 1)
|
|
final_down = torch.cat([f_down, w_down_scaled], dim=0)
|
|
|
|
w_up_scaled = w_up_train * s1.view(1, -1, 1, 1)
|
|
final_up = torch.cat([w_up_scaled, f_up], dim=1)
|
|
|
|
else:
|
|
continue
|
|
|
|
reconstructed_lora[comfy_key_base + ".lora_down.weight"] = final_down.cpu()
|
|
reconstructed_lora[comfy_key_base + ".lora_up.weight"] = final_up.cpu()
|
|
count += 1
|
|
|
|
except Exception as e:
|
|
print(f"Error reconstructing Null-LaLoRA {base_name}: {e}")
|
|
|
|
print(f"DEBUG: Null-LaLoRA Applied Layers: {count} (DoRA: {use_dora})")
|
|
if count == 0: return (model, clip)
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, reconstructed_lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UniversalLoRALoader": UniversalLoRALoader,
|
|
"LoadUniversalWeights": LoadUniversalWeights,
|
|
"ApplyUniversalWeights": ApplyUniversalWeights,
|
|
"ApplyUniversalLoRA": ApplyUniversalLoRA,
|
|
"ApplyNullLaLoRA": ApplyNullLaLoRA
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UniversalLoRALoader": "Load Universal Subspace (Basis)",
|
|
"LoadUniversalWeights": "Load Universal Weights (Alpha)",
|
|
"ApplyUniversalWeights": "Apply Universal Weights",
|
|
"ApplyUniversalLoRA": "Approximating LoRA (Project Target)",
|
|
"ApplyNullLaLoRA": "Apply Null-LaLoRA (Universal + Null)"
|
|
} |