kohya_ss/tools/extract_lora_from_models-new.py

import sys
import os
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)