mirror of https://github.com/bmaltais/kohya_ss
563 lines
28 KiB
Python
563 lines
28 KiB
Python
import sys
|
|
import os
|
|
|
|
# 1. Add sd-scripts directory to sys.path
|
|
# This block can now be potentially removed if no other sd-scripts imports are needed
|
|
# OR kept if there's a chance of re-introducing some utilities for other purposes.
|
|
# For full removal of the sd-scripts dependency for *this script's execution*,
|
|
# ensure no other `from library...` or `from networks...` exist.
|
|
# script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# project_root = os.path.dirname(script_dir)
|
|
# sd_scripts_dir_path = os.path.join(project_root, "sd-scripts")
|
|
|
|
# if sd_scripts_dir_path not in sys.path:
|
|
# sys.path.insert(0, sd_scripts_dir_path)
|
|
|
|
# Now you can import from the library package and the networks package
|
|
# try:
|
|
# # model_util and sdxl_model_util REMOVED from here
|
|
# # from library.utils import setup_logging # REMOVED
|
|
# # from networks import lora # REMOVED
|
|
# except ImportError as e:
|
|
# print(f"Error importing from sd-scripts. Please check your sd-scripts folder structure.")
|
|
# # print(f"Attempted to load from: {sd_scripts_dir_path}") # If path addition is removed
|
|
# print(f"Original error: {e}")
|
|
# print("Current sys.path relevant entries:")
|
|
# for p in sys.path:
|
|
# if "sd-scripts" in p or "kohya_ss" in p: # Adjust if sd_scripts_dir_path is removed
|
|
# print(p)
|
|
# raise
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
from tqdm import tqdm
|
|
import logging # Import for logging
|
|
|
|
# NEW: Add diffusers import for model loading
|
|
try:
|
|
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
|
|
except ImportError:
|
|
print("Diffusers library not found. Please install it: pip install diffusers transformers accelerate")
|
|
raise
|
|
|
|
# --- Localized Logging Setup ---
|
|
def _local_setup_logging(log_level=logging.INFO):
|
|
"""
|
|
Sets up basic logging to console.
|
|
"""
|
|
logging.basicConfig(
|
|
level=log_level,
|
|
format="%(asctime)s %(levelname)-8s %(name)s - %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
_local_setup_logging() # Initialize logging
|
|
logger = logging.getLogger(__name__) # Get logger for this module
|
|
|
|
MIN_SV = 1e-6
|
|
|
|
# --- Localized sd-scripts constants and utility functions ---
|
|
_LOCAL_MODEL_VERSION_SDXL_BASE_V1_0 = "sdxl_v10"
|
|
|
|
def _local_get_model_version_str_for_sd1_sd2(is_v2: bool, is_v_parameterization: bool) -> str:
|
|
if is_v2:
|
|
return "v2-v" if is_v_parameterization else "v2"
|
|
return "v1"
|
|
|
|
# --- Localized LoRA Placeholder and Network Creation ---
|
|
class LocalLoRAModulePlaceholder:
|
|
def __init__(self, lora_name: str, org_module: torch.nn.Module):
|
|
self.lora_name = lora_name
|
|
self.org_module = org_module
|
|
# Add other attributes if _calculate_module_diffs_and_check needs them,
|
|
# but it primarily uses .lora_name and .org_module.weight
|
|
|
|
def _local_create_network_placeholders(text_encoders: list, unet: torch.nn.Module, lora_conv_dim_init: int):
|
|
"""
|
|
Creates placeholders for LoRA-able modules in text encoders and UNet.
|
|
Mimics the module identification and naming of sd-scripts' lora.create_network.
|
|
`lora_conv_dim_init`: If > 0, Conv2d layers are considered for LoRA.
|
|
"""
|
|
unet_loras = []
|
|
text_encoder_loras = []
|
|
|
|
# Target U-Net modules
|
|
for name, module in unet.named_modules():
|
|
lora_name = "lora_unet_" + name.replace(".", "_")
|
|
if isinstance(module, torch.nn.Linear):
|
|
unet_loras.append(LocalLoRAModulePlaceholder(lora_name, module))
|
|
elif isinstance(module, torch.nn.Conv2d):
|
|
if lora_conv_dim_init > 0: # Only consider conv layers if conv_dim > 0
|
|
# Kernel size check might be relevant if sd-scripts has specific logic,
|
|
# but for diffing, any conv is a candidate if conv_dim > 0.
|
|
# SVD will later handle rank based on actual layer type (1x1 vs 3x3).
|
|
unet_loras.append(LocalLoRAModulePlaceholder(lora_name, module))
|
|
|
|
# Target Text Encoder modules
|
|
for i, text_encoder in enumerate(text_encoders):
|
|
if text_encoder is None: # SDXL can have None TEs if not loaded
|
|
continue
|
|
# Determine prefix based on number of text encoders (for SDXL compatibility)
|
|
te_prefix = f"lora_te{i+1}_" if len(text_encoders) > 1 else "lora_te_"
|
|
|
|
for name, module in text_encoder.named_modules():
|
|
lora_name = te_prefix + name.replace(".", "_")
|
|
if isinstance(module, torch.nn.Linear):
|
|
text_encoder_loras.append(LocalLoRAModulePlaceholder(lora_name, module))
|
|
# Conv2d in text encoders is rare but check just in case (sd-scripts might)
|
|
elif isinstance(module, torch.nn.Conv2d):
|
|
if lora_conv_dim_init > 0:
|
|
text_encoder_loras.append(LocalLoRAModulePlaceholder(lora_name, module))
|
|
|
|
logger.info(f"Found {len(text_encoder_loras)} LoRA-able placeholder modules in Text Encoders.")
|
|
logger.info(f"Found {len(unet_loras)} LoRA-able placeholder modules in U-Net.")
|
|
return text_encoder_loras, unet_loras
|
|
|
|
|
|
# --- Singular Value Indexing Functions (Unchanged) ---
|
|
def index_sv_cumulative(S, target):
|
|
original_sum = float(torch.sum(S))
|
|
cumulative_sums = torch.cumsum(S, dim=0) / original_sum
|
|
index = int(torch.searchsorted(cumulative_sums, target)) + 1
|
|
index = max(1, min(index, len(S) - 1))
|
|
return index
|
|
|
|
def index_sv_fro(S, target):
|
|
S_squared = S.pow(2)
|
|
S_fro_sq = float(torch.sum(S_squared))
|
|
sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq
|
|
index = int(torch.searchsorted(sum_S_squared, target**2)) + 1
|
|
index = max(1, min(index, len(S) - 1))
|
|
return index
|
|
|
|
def index_sv_ratio(S, target):
|
|
max_sv = S[0]
|
|
min_sv = max_sv / target
|
|
index = int(torch.sum(S > min_sv).item())
|
|
index = max(1, min(index, len(S) - 1))
|
|
return index
|
|
|
|
def index_sv_knee(S, MIN_SV_KNEE=1e-8):
|
|
n = len(S)
|
|
if n < 3: return 1
|
|
s_max, s_min = S[0], S[-1]
|
|
if s_max - s_min < MIN_SV_KNEE: return 1
|
|
s_normalized = (S - s_min) / (s_max - s_min)
|
|
x_normalized = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
|
|
distances = (x_normalized + s_normalized - 1).abs()
|
|
knee_index_0based = torch.argmax(distances).item()
|
|
rank = knee_index_0based + 1
|
|
rank = max(1, min(rank, n - 1))
|
|
return rank
|
|
|
|
def index_sv_cumulative_knee(S, min_sv_threshold=1e-8):
|
|
n = len(S)
|
|
if n < 3: return 1
|
|
s_sum = torch.sum(S)
|
|
if s_sum < min_sv_threshold: return 1
|
|
y_values = torch.cumsum(S, dim=0) / s_sum
|
|
y_min, y_max = y_values[0], y_values[n-1]
|
|
if y_max - y_min < min_sv_threshold: return 1
|
|
y_norm = (y_values - y_min) / (y_max - y_min)
|
|
x_norm = torch.linspace(0, 1, n, device=S.device, dtype=S.dtype)
|
|
distances = (y_norm - x_norm).abs()
|
|
knee_index_0based = torch.argmax(distances).item()
|
|
rank = knee_index_0based + 1
|
|
rank = max(1, min(rank, n - 1))
|
|
return rank
|
|
|
|
def index_sv_rel_decrease(S, tau=0.1):
|
|
if len(S) < 2: return 1
|
|
ratios = S[1:] / S[:-1]
|
|
for k in range(len(ratios)):
|
|
if ratios[k] < tau:
|
|
return k + 1
|
|
return len(S)
|
|
|
|
# --- Utility Functions ---
|
|
def _str_to_dtype(p):
|
|
if p == "float": return torch.float
|
|
if p == "fp16": return torch.float16
|
|
if p == "bf16": return torch.bfloat16
|
|
return None
|
|
|
|
def save_to_file(file_name, state_dict_to_save, dtype, metadata=None):
|
|
state_dict_final = {}
|
|
for key, value in state_dict_to_save.items():
|
|
if isinstance(value, torch.Tensor) and dtype is not None:
|
|
state_dict_final[key] = value.to(dtype)
|
|
else:
|
|
state_dict_final[key] = value
|
|
|
|
if os.path.splitext(file_name)[1] == ".safetensors":
|
|
save_file(state_dict_final, file_name, metadata=metadata)
|
|
else:
|
|
torch.save(state_dict_final, file_name)
|
|
|
|
def _build_local_sai_metadata(title, creation_time, is_v2_flag, is_v_param_flag, is_sdxl_flag):
|
|
metadata = {}
|
|
metadata["ss_sd_model_name"] = str(title)
|
|
metadata["ss_creation_time"] = str(int(creation_time))
|
|
if is_sdxl_flag:
|
|
metadata["ss_base_model_version"] = "sdxl_v10"
|
|
metadata["ss_sdxl_model_version"] = "1.0"
|
|
if is_v_param_flag:
|
|
metadata["ss_v_parameterization"] = "true"
|
|
elif is_v2_flag:
|
|
metadata["ss_base_model_version"] = "sd_v2"
|
|
if is_v_param_flag:
|
|
metadata["ss_v_parameterization"] = "true"
|
|
else:
|
|
metadata["ss_base_model_version"] = "sd_v1"
|
|
if is_v_param_flag:
|
|
metadata["ss_v_parameterization"] = "true"
|
|
return metadata
|
|
|
|
# --- MODIFIED Helper Functions for Model Loading ---
|
|
def _load_sd_model_components(model_path, is_v2_flag, target_device_override, load_dtype_torch):
|
|
logger.info(f"Loading SD model using Diffusers.StableDiffusionPipeline from: {model_path}")
|
|
pipeline = StableDiffusionPipeline.from_single_file(
|
|
model_path,
|
|
torch_dtype=load_dtype_torch
|
|
)
|
|
eff_device = target_device_override if target_device_override else "cpu"
|
|
text_encoder = pipeline.text_encoder.to(eff_device)
|
|
unet = pipeline.unet.to(eff_device)
|
|
text_encoders = [text_encoder]
|
|
logger.info(f"Loaded SD model components. UNet device: {unet.device}, TextEncoder device: {text_encoder.device}")
|
|
return text_encoders, unet
|
|
|
|
def _load_sdxl_model_components(model_path, target_device_override, load_dtype_torch):
|
|
actual_load_device = target_device_override if target_device_override else "cpu"
|
|
logger.info(f"Loading SDXL model using Diffusers.StableDiffusionXLPipeline from: {model_path} to device: {actual_load_device}")
|
|
pipeline = StableDiffusionXLPipeline.from_single_file(
|
|
model_path,
|
|
torch_dtype=load_dtype_torch
|
|
)
|
|
pipeline.to(actual_load_device)
|
|
text_encoder = pipeline.text_encoder
|
|
text_encoder_2 = pipeline.text_encoder_2
|
|
unet = pipeline.unet
|
|
text_encoders = [text_encoder, text_encoder_2]
|
|
logger.info(f"Loaded SDXL model components. UNet device: {unet.device}, TextEncoder1 device: {text_encoder.device}, TextEncoder2 device: {text_encoder_2.device}")
|
|
return text_encoders, unet
|
|
|
|
def _calculate_module_diffs_and_check(module_loras_o, module_loras_t, diff_calc_device, min_diff_thresh, module_type_str):
|
|
diffs_map = {}
|
|
is_different_flag = False
|
|
first_diff_logged = False
|
|
for lora_o, lora_t in zip(module_loras_o, module_loras_t):
|
|
lora_name = lora_o.lora_name
|
|
if lora_o.org_module is None or lora_t.org_module is None or \
|
|
not hasattr(lora_o.org_module, 'weight') or lora_o.org_module.weight is None or \
|
|
not hasattr(lora_t.org_module, 'weight') or lora_t.org_module.weight is None:
|
|
logger.warning(f"Skipping {lora_name} in {module_type_str} due to missing org_module or weight.")
|
|
continue
|
|
weight_o = lora_o.org_module.weight
|
|
weight_t = lora_t.org_module.weight
|
|
if str(weight_o.device) != str(diff_calc_device): weight_o = weight_o.to(diff_calc_device)
|
|
if str(weight_t.device) != str(diff_calc_device): weight_t = weight_t.to(diff_calc_device)
|
|
diff = weight_t - weight_o
|
|
diffs_map[lora_name] = diff
|
|
current_max_diff = torch.max(torch.abs(diff))
|
|
if not is_different_flag and current_max_diff > min_diff_thresh:
|
|
is_different_flag = True
|
|
if not first_diff_logged:
|
|
logger.info(f"{module_type_str} '{lora_name}' differs: max diff {current_max_diff} > {min_diff_thresh}")
|
|
first_diff_logged = True
|
|
return diffs_map, is_different_flag
|
|
|
|
def _determine_rank(S_values, dynamic_method_name, dynamic_param_value, max_rank_limit,
|
|
module_eff_in_dim, module_eff_out_dim, min_sv_threshold=MIN_SV):
|
|
if not S_values.numel() or S_values[0] <= min_sv_threshold: return 1
|
|
rank = 0
|
|
if dynamic_method_name == "sv_ratio": rank = index_sv_ratio(S_values, dynamic_param_value)
|
|
elif dynamic_method_name == "sv_cumulative": rank = index_sv_cumulative(S_values, dynamic_param_value)
|
|
elif dynamic_method_name == "sv_fro": rank = index_sv_fro(S_values, dynamic_param_value)
|
|
elif dynamic_method_name == "sv_knee": rank = index_sv_knee(S_values, min_sv_threshold)
|
|
elif dynamic_method_name == "sv_cumulative_knee": rank = index_sv_cumulative_knee(S_values, min_sv_threshold)
|
|
elif dynamic_method_name == "sv_rel_decrease": rank = index_sv_rel_decrease(S_values, dynamic_param_value)
|
|
else: rank = max_rank_limit
|
|
rank = min(rank, max_rank_limit, module_eff_in_dim, module_eff_out_dim, len(S_values))
|
|
rank = max(1, rank)
|
|
return rank
|
|
|
|
def _construct_lora_weights_from_svd_components(U_full, S_all_values, Vh_full, rank,
|
|
clamp_quantile_val, is_conv2d, is_conv2d_3x3,
|
|
conv_kernel_size,
|
|
module_out_channels, module_in_channels,
|
|
target_device_for_final_weights, target_dtype_for_final_weights):
|
|
S_k = S_all_values[:rank]
|
|
U_k = U_full[:, :rank]
|
|
Vh_k = Vh_full[:rank, :]
|
|
S_k_non_negative = torch.clamp(S_k, min=0.0)
|
|
s_sqrt = torch.sqrt(S_k_non_negative)
|
|
U_final = U_k * s_sqrt.unsqueeze(0)
|
|
Vh_final = Vh_k * s_sqrt.unsqueeze(1)
|
|
dist = torch.cat([U_final.flatten(), Vh_final.flatten()])
|
|
hi_val = torch.quantile(dist, clamp_quantile_val)
|
|
if hi_val == 0 and torch.max(torch.abs(dist)) > 1e-9:
|
|
logger.debug(f"Clamping hi_val is zero for non-zero distribution. Max abs val: {torch.max(torch.abs(dist))}. Quantile: {clamp_quantile_val}")
|
|
U_clamped = U_final.clamp(-hi_val, hi_val)
|
|
Vh_clamped = Vh_final.clamp(-hi_val, hi_val)
|
|
if is_conv2d:
|
|
U_clamped = U_clamped.reshape(module_out_channels, rank, 1, 1)
|
|
if is_conv2d_3x3:
|
|
Vh_clamped = Vh_clamped.reshape(rank, module_in_channels, *conv_kernel_size)
|
|
else:
|
|
Vh_clamped = Vh_clamped.reshape(rank, module_in_channels, 1, 1)
|
|
U_clamped = U_clamped.to(target_device_for_final_weights, dtype=target_dtype_for_final_weights).contiguous()
|
|
Vh_clamped = Vh_clamped.to(target_device_for_final_weights, dtype=target_dtype_for_final_weights).contiguous()
|
|
return U_clamped, Vh_clamped
|
|
|
|
def _log_svd_stats(lora_module_name, S_all_values, rank_used, min_sv_for_calc=MIN_SV):
|
|
if not S_all_values.numel():
|
|
logger.info(f"{lora_module_name:75} | rank: {rank_used}, SVD not performed (empty singular values).")
|
|
return
|
|
S_cpu = S_all_values.to('cpu')
|
|
s_sum_total = float(torch.sum(S_cpu))
|
|
s_sum_rank = float(torch.sum(S_cpu[:rank_used]))
|
|
fro_orig_total = float(torch.sqrt(torch.sum(S_cpu.pow(2))))
|
|
fro_reconstructed_rank = float(torch.sqrt(torch.sum(S_cpu[:rank_used].pow(2))))
|
|
ratio_sv = float('inf')
|
|
if rank_used > 0 and S_cpu[rank_used - 1].abs() > min_sv_for_calc:
|
|
ratio_sv = S_cpu[0] / S_cpu[rank_used - 1]
|
|
sum_s_retained_percentage = (s_sum_rank / s_sum_total) if s_sum_total > min_sv_for_calc else 1.0
|
|
fro_retained_percentage = (fro_reconstructed_rank / fro_orig_total) if fro_orig_total > min_sv_for_calc else 1.0
|
|
logger.info(
|
|
f"{lora_module_name:75} | rank: {rank_used}, "
|
|
f"sum(S) retained: {sum_s_retained_percentage:.2%}, "
|
|
f"Frobenius norm retained: {fro_retained_percentage:.2%}, "
|
|
f"max_retained_sv/min_retained_sv ratio: {ratio_sv:.2f}"
|
|
)
|
|
|
|
def _prepare_lora_metadata(output_path, is_v2_flag, kohya_base_model_version_str, network_conv_dim_val,
|
|
use_dynamic_method_flag, network_dim_config_val,
|
|
is_v_param_flag, is_sdxl_flag, skip_sai_meta):
|
|
net_kwargs = {"conv_dim": str(network_conv_dim_val), "conv_alpha": str(float(network_conv_dim_val))} if network_conv_dim_val is not None else {}
|
|
if use_dynamic_method_flag:
|
|
network_dim_meta = "Dynamic"
|
|
network_alpha_meta = "Dynamic"
|
|
else:
|
|
network_dim_meta = str(network_dim_config_val)
|
|
network_alpha_meta = str(float(network_dim_config_val))
|
|
final_metadata = {
|
|
"ss_v2": str(is_v2_flag),
|
|
"ss_base_model_version": kohya_base_model_version_str,
|
|
"ss_network_module": "networks.lora", # This remains for compatibility with tools expecting it
|
|
"ss_network_dim": network_dim_meta,
|
|
"ss_network_alpha": network_alpha_meta,
|
|
"ss_network_args": json.dumps(net_kwargs),
|
|
"ss_lowram": "False",
|
|
"ss_num_train_images": "N/A",
|
|
}
|
|
if not skip_sai_meta:
|
|
title = os.path.splitext(os.path.basename(output_path))[0]
|
|
current_time = time.time()
|
|
sai_metadata_content = _build_local_sai_metadata(
|
|
title=title, creation_time=current_time, is_v2_flag=is_v2_flag,
|
|
is_v_param_flag=is_v_param_flag, is_sdxl_flag=is_sdxl_flag
|
|
)
|
|
final_metadata.update(sai_metadata_content)
|
|
return final_metadata
|
|
|
|
# --- Main SVD Function ---
|
|
def svd(
|
|
model_org=None, model_tuned=None, save_to=None, dim=4, v2=None, sdxl=None,
|
|
conv_dim=None, v_parameterization=None, device=None, save_precision=None,
|
|
clamp_quantile=0.99, min_diff=0.01, no_metadata=False, load_precision=None,
|
|
load_original_model_to=None, load_tuned_model_to=None,
|
|
dynamic_method=None, dynamic_param=None, verbose=False,
|
|
):
|
|
actual_v_parameterization = v2 if v_parameterization is None else v_parameterization
|
|
load_dtype_torch = _str_to_dtype(load_precision)
|
|
save_dtype_torch = _str_to_dtype(save_precision) if save_precision else torch.float
|
|
|
|
svd_computation_device = torch.device(device if device else "cuda" if torch.cuda.is_available() else "cpu")
|
|
logger.info(f"Using SVD computation device: {svd_computation_device}")
|
|
diff_calculation_device = torch.device("cpu")
|
|
logger.info(f"Calculating weight differences on: {diff_calculation_device}")
|
|
final_weights_device = torch.device("cpu")
|
|
|
|
if not sdxl:
|
|
text_encoders_o, unet_o = _load_sd_model_components(model_org, v2, load_original_model_to, load_dtype_torch)
|
|
text_encoders_t, unet_t = _load_sd_model_components(model_tuned, v2, load_tuned_model_to, load_dtype_torch)
|
|
kohya_model_version = _local_get_model_version_str_for_sd1_sd2(v2, actual_v_parameterization)
|
|
else:
|
|
text_encoders_o, unet_o = _load_sdxl_model_components(model_org, load_original_model_to, load_dtype_torch)
|
|
text_encoders_t, unet_t = _load_sdxl_model_components(model_tuned, load_tuned_model_to, load_dtype_torch)
|
|
kohya_model_version = _LOCAL_MODEL_VERSION_SDXL_BASE_V1_0
|
|
|
|
# Determine lora_conv_dim_init based on conv_dim argument for network creation
|
|
# The original script used init_dim_val (1) if conv_dim was None.
|
|
# Here, conv_dim is already defaulted to args.dim if None by the main block.
|
|
# So, lora_conv_dim_init will be args.conv_dim (which defaults to args.dim).
|
|
# If args.conv_dim was explicitly 0, this would be 0.
|
|
lora_conv_dim_init_val = conv_dim # conv_dim is args.conv_dim (or args.dim)
|
|
|
|
# Create LoRA placeholders using the localized function
|
|
text_encoder_loras_o, unet_loras_o = _local_create_network_placeholders(text_encoders_o, unet_o, lora_conv_dim_init_val)
|
|
text_encoder_loras_t, unet_loras_t = _local_create_network_placeholders(text_encoders_t, unet_t, lora_conv_dim_init_val) # same conv_dim logic for tuned
|
|
|
|
# Group LoRA placeholders for easier processing (mimicking LoraNetwork structure somewhat)
|
|
class LocalLoraNetworkPlaceholder:
|
|
def __init__(self, te_loras, unet_loras_list):
|
|
self.text_encoder_loras = te_loras
|
|
self.unet_loras = unet_loras_list
|
|
|
|
lora_network_o = LocalLoraNetworkPlaceholder(text_encoder_loras_o, unet_loras_o)
|
|
lora_network_t = LocalLoraNetworkPlaceholder(text_encoder_loras_t, unet_loras_t)
|
|
|
|
assert len(lora_network_o.text_encoder_loras) == len(lora_network_t.text_encoder_loras), \
|
|
f"Model versions (based on identified LoRA-able TE modules) differ: {len(lora_network_o.text_encoder_loras)} vs {len(lora_network_t.text_encoder_loras)} TEs"
|
|
|
|
all_diffs = {}
|
|
te_diffs, text_encoder_different = _calculate_module_diffs_and_check(
|
|
lora_network_o.text_encoder_loras, lora_network_t.text_encoder_loras,
|
|
diff_calculation_device, min_diff, "Text Encoder"
|
|
)
|
|
|
|
if text_encoder_different:
|
|
all_diffs.update(te_diffs)
|
|
else:
|
|
logger.warning("Text encoders are considered identical based on min_diff. Not extracting TE LoRA.")
|
|
# To prevent processing empty list later, ensure it's empty if no diffs
|
|
lora_network_o.text_encoder_loras = []
|
|
del text_encoders_t # Free memory early
|
|
|
|
unet_diffs, _ = _calculate_module_diffs_and_check(
|
|
lora_network_o.unet_loras, lora_network_t.unet_loras,
|
|
diff_calculation_device, min_diff, "U-Net"
|
|
)
|
|
all_diffs.update(unet_diffs)
|
|
del lora_network_t, unet_t # Free memory early
|
|
|
|
# Ensure lora_names_to_process only includes modules from lora_network_o
|
|
# that are actually present (e.g., if TEs were skipped)
|
|
lora_names_to_process = set()
|
|
if text_encoder_different: # Only add TE loras if they were deemed different
|
|
lora_names_to_process.update(p.lora_name for p in lora_network_o.text_encoder_loras)
|
|
lora_names_to_process.update(p.lora_name for p in lora_network_o.unet_loras)
|
|
|
|
logger.info("Extracting and resizing LoRA via SVD")
|
|
lora_weights = {}
|
|
with torch.no_grad():
|
|
for lora_name in tqdm(lora_names_to_process):
|
|
if lora_name not in all_diffs:
|
|
logger.warning(f"Skipping {lora_name} as no diff was calculated for it (e.g., Text Encoders were identical).")
|
|
continue
|
|
original_diff_tensor = all_diffs[lora_name]
|
|
is_conv2d_layer = len(original_diff_tensor.size()) == 4
|
|
kernel_s = original_diff_tensor.size()[2:4] if is_conv2d_layer else None
|
|
is_conv2d_3x3_layer = is_conv2d_layer and kernel_s != (1, 1)
|
|
module_true_out_channels, module_true_in_channels = original_diff_tensor.size()[0:2]
|
|
mat_for_svd = original_diff_tensor.to(svd_computation_device, dtype=torch.float)
|
|
if is_conv2d_layer:
|
|
if is_conv2d_3x3_layer: mat_for_svd = mat_for_svd.flatten(start_dim=1)
|
|
else: mat_for_svd = mat_for_svd.squeeze()
|
|
if mat_for_svd.numel() == 0 or mat_for_svd.shape[0] == 0 or mat_for_svd.shape[1] == 0 :
|
|
logger.warning(f"Skipping SVD for {lora_name} due to empty/invalid shape: {mat_for_svd.shape}")
|
|
continue
|
|
try:
|
|
U_full, S_full, Vh_full = torch.linalg.svd(mat_for_svd)
|
|
except Exception as e:
|
|
logger.error(f"SVD failed for {lora_name} with shape {mat_for_svd.shape}. Error: {e}")
|
|
continue
|
|
|
|
# Max rank for SVD is based on 'dim' for linear and 'conv_dim' for conv3x3
|
|
# The original `current_max_rank` logic was:
|
|
# current_max_rank = dim if not is_conv2d_3x3_layer or conv_dim is None else conv_dim
|
|
# Here, `dim` is args.dim and `conv_dim` is args.conv_dim (defaulted to args.dim)
|
|
module_specific_max_rank = conv_dim if is_conv2d_3x3_layer else dim
|
|
|
|
eff_out_dim, eff_in_dim = mat_for_svd.shape[0], mat_for_svd.shape[1]
|
|
rank = _determine_rank(S_full, dynamic_method, dynamic_param,
|
|
module_specific_max_rank, eff_in_dim, eff_out_dim, MIN_SV)
|
|
U_clamped, Vh_clamped = _construct_lora_weights_from_svd_components(
|
|
U_full, S_full, Vh_full, rank, clamp_quantile,
|
|
is_conv2d_layer, is_conv2d_3x3_layer, kernel_s,
|
|
module_true_out_channels, module_true_in_channels,
|
|
final_weights_device, save_dtype_torch
|
|
)
|
|
lora_weights[lora_name] = (U_clamped, Vh_clamped)
|
|
if verbose: _log_svd_stats(lora_name, S_full, rank, MIN_SV)
|
|
|
|
lora_sd = {}
|
|
for lora_name, (up_weight, down_weight) in lora_weights.items():
|
|
lora_sd[lora_name + ".lora_up.weight"] = up_weight
|
|
lora_sd[lora_name + ".lora_down.weight"] = down_weight
|
|
# Alpha is set to the rank (dim of down_weight's 0th axis, which is rank)
|
|
lora_sd[lora_name + ".alpha"] = torch.tensor(down_weight.size()[0], dtype=save_dtype_torch, device=final_weights_device)
|
|
|
|
del text_encoders_o, unet_o, lora_network_o, all_diffs # Clean up original models and placeholders
|
|
if 'torch' in sys.modules and hasattr(torch, 'cuda') and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
if not os.path.exists(os.path.dirname(save_to)) and os.path.dirname(save_to) != "":
|
|
os.makedirs(os.path.dirname(save_to), exist_ok=True)
|
|
|
|
metadata_to_save = _prepare_lora_metadata(
|
|
output_path=save_to,
|
|
is_v2_flag=v2,
|
|
kohya_base_model_version_str=kohya_model_version,
|
|
network_conv_dim_val=conv_dim, # This is args.conv_dim (defaulted to args.dim)
|
|
use_dynamic_method_flag=bool(dynamic_method),
|
|
network_dim_config_val=dim, # This is args.dim
|
|
is_v_param_flag=actual_v_parameterization,
|
|
is_sdxl_flag=sdxl,
|
|
skip_sai_meta=no_metadata
|
|
)
|
|
|
|
save_to_file(save_to, lora_sd, save_dtype_torch, metadata_to_save)
|
|
logger.info(f"LoRA saved to: {save_to}")
|
|
|
|
def setup_parser():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--v2", action="store_true", help="Load Stable Diffusion v2.x model")
|
|
parser.add_argument("--v_parameterization", action="store_true", help="Set v-parameterization metadata (defaults to v2 if --v2 is set)")
|
|
parser.add_argument("--sdxl", action="store_true", help="Load Stable Diffusion SDXL base model")
|
|
parser.add_argument("--load_precision", type=str, choices=["float", "fp16", "bf16"], default=None, help="Precision for loading models (applied after initial load)")
|
|
parser.add_argument("--save_precision", type=str, choices=["float", "fp16", "bf16"], default="float", help="Precision for saving LoRA weights")
|
|
parser.add_argument("--model_org", type=str, required=True, help="Original Stable Diffusion model (ckpt/safetensors)")
|
|
parser.add_argument("--model_tuned", type=str, required=True, help="Tuned Stable Diffusion model (ckpt/safetensors)")
|
|
parser.add_argument("--save_to", type=str, required=True, help="Output file name (ckpt/safetensors)")
|
|
parser.add_argument("--dim", type=int, default=4, help="Max dimension (rank) of LoRA for linear layers")
|
|
parser.add_argument("--conv_dim", type=int, default=None, help="Max dimension (rank) of LoRA for Conv2d-3x3. Defaults to 'dim' if not set.")
|
|
parser.add_argument("--device", type=str, default=None, help="Device for SVD computation (e.g., cuda, cpu). Defaults to cuda if available, else cpu.")
|
|
parser.add_argument("--clamp_quantile", type=float, default=0.99, help="Quantile for clamping weights")
|
|
parser.add_argument("--min_diff", type=float, default=0.01, help="Minimum weight difference to extract LoRA for a module")
|
|
parser.add_argument("--no_metadata", action="store_true", help="Omit detailed metadata from SAI and Kohya_ss")
|
|
parser.add_argument("--load_original_model_to", type=str, default=None, help="Device for original model (e.g. 'cpu', 'cuda:0'). Defaults to CPU for SD1/2, honored for SDXL.")
|
|
parser.add_argument("--load_tuned_model_to", type=str, default=None, help="Device for tuned model (e.g. 'cpu', 'cuda:0'). Defaults to CPU for SD1/2, honored for SDXL.")
|
|
parser.add_argument("--dynamic_param", type=float, help="Parameter for dynamic rank reduction")
|
|
parser.add_argument("--verbose", action="store_true", help="Show detailed rank reduction info for each module")
|
|
parser.add_argument(
|
|
"--dynamic_method", type=str,
|
|
choices=[None, "sv_ratio", "sv_fro", "sv_cumulative", "sv_knee", "sv_rel_decrease", "sv_cumulative_knee"],
|
|
default=None, help="Dynamic rank reduction method"
|
|
)
|
|
return parser
|
|
|
|
if __name__ == "__main__":
|
|
parser = setup_parser()
|
|
args = parser.parse_args()
|
|
|
|
if args.conv_dim is None:
|
|
args.conv_dim = args.dim # Default conv_dim to dim if not provided
|
|
logger.info(f"--conv_dim not set, using value of --dim: {args.conv_dim}")
|
|
|
|
methods_requiring_param = ["sv_ratio", "sv_fro", "sv_cumulative", "sv_rel_decrease"]
|
|
if args.dynamic_method in methods_requiring_param and args.dynamic_param is None:
|
|
parser.error(f"Dynamic method '{args.dynamic_method}' requires --dynamic_param to be set.")
|
|
|
|
if not args.dynamic_method: # Ranks must be positive if not using dynamic method
|
|
if args.dim <= 0: parser.error(f"--dim (rank) must be > 0. Got {args.dim}")
|
|
if args.conv_dim <=0: parser.error(f"--conv_dim (rank) must be > 0. Got {args.conv_dim}") # Check after defaulting
|
|
|
|
if MIN_SV <= 0: logger.warning(f"Global MIN_SV ({MIN_SV}) should be positive.")
|
|
|
|
svd_args = vars(args).copy()
|
|
svd(**svd_args) |