Modify AI-Toolkit Z Image Turbo Lora

modify_zlora.py

import argparse
import re
import torch
from safetensors.torch import load_file, save_file
import sys
import os


# ============================================================================
# 輔助函數：鍵和層處理
# ============================================================================

def extract_layer_number(key):
    """從 key 中提取 layer 編號，例如 'layers.0.xxx' -> 0，'layers.29.xxx' -> 29"""
    match = re.search(r'layers\.(\d+)', key)
    if match:
        return int(match.group(1))
    return None


def is_lora_key(key):
    """檢查 key 是否為 LoRA 鍵（lora_A.weight 或 lora_B.weight）"""
    return 'lora_A.weight' in key or 'lora_B.weight' in key


def get_lora_rank_info(key, shape):
    """
    獲取 LoRA 鍵的 rank 維度資訊
    返回: (rank_dim_index, current_rank) 或 None（如果不是 LoRA 鍵）
    """
    if 'lora_A.weight' in key:
        # lora_A.weight 的 rank 是第一個維度
        return (0, shape[0])
    elif 'lora_B.weight' in key:
        # lora_B.weight 的 rank 是最後一個維度
        return (-1, shape[-1])
    return None


# ============================================================================
# 輔助函數：張量對齊和 rank 調整
# ============================================================================

def align_tensor_shape(tensor, target_shape, dtype, device):
    """將張量調整到目標形狀，不足的部分用零填充"""
    if tensor.shape == target_shape:
        return tensor

    # 檢查維度數量是否相同
    if len(tensor.shape) != len(target_shape):
        raise ValueError(
            f"無法對齊張量：維度數量不相符 "
            f"({len(tensor.shape)}D vs {len(target_shape)}D). "
            f"形狀：{tensor.shape} vs {target_shape}"
        )

    # 建立目標形狀的零張量
    aligned = torch.zeros(target_shape, dtype=dtype, device=device)

    # 構建切片，將原始張量的數據複製到對齊後的張量左上角
    # 對於每個維度，取較小的值作為複製範圍
    slices = tuple(slice(0, min(tensor.shape[i], target_shape[i])) for i in range(len(target_shape)))

    # 將原始張量的數據複製到對齊後的張量中
    aligned[slices] = tensor[slices]

    return aligned


def reduce_rank_average(tensor, rank_dim_index, current_rank, target_rank, dtype, device):
    """
    使用平均聚合方法將 rank 從 current_rank 降到 target_rank
    將多個維度平均成 target_rank 個維度
    """
    if rank_dim_index == 0:
        # lora_A: (current_rank, ...) -> (target_rank, ...)
        # 計算每個目標維度需要聚合多少個原始維度
        group_size = current_rank / target_rank
        aligned = torch.zeros((target_rank,) + tensor.shape[1:], dtype=dtype, device=device)

        for i in range(target_rank):
            start_idx = int(i * group_size)
            end_idx = int((i + 1) * group_size) if i < target_rank - 1 else current_rank
            # 對這一段維度求平均
            aligned[i, ...] = tensor[start_idx:end_idx, ...].mean(dim=0)
    else:
        # lora_B: (..., current_rank) -> (..., target_rank)
        group_size = current_rank / target_rank
        aligned = torch.zeros(tensor.shape[:-1] + (target_rank,), dtype=dtype, device=device)

        for i in range(target_rank):
            start_idx = int(i * group_size)
            end_idx = int((i + 1) * group_size) if i < target_rank - 1 else current_rank
            # 對這一段維度求平均
            aligned[..., i] = tensor[..., start_idx:end_idx].mean(dim=-1)

    return aligned


def reduce_rank_svd(tensor, rank_dim_index, current_rank, target_rank, dtype, device):
    """
    使用 SVD 進行低秩近似，將 rank 從 current_rank 降到 target_rank
    注意：SVD 需要 float32 精度，所以先轉換為 float32，計算完成後再轉回原始類型
    """
    # SVD 需要 float32 精度，先轉換
    tensor_fp32 = tensor.float()

    if rank_dim_index == 0:
        # lora_A: (current_rank, in_features)
        # 轉置為 (in_features, current_rank) 進行 SVD
        tensor_t = tensor_fp32.t()  # (in_features, current_rank)
        U, S, Vh = torch.linalg.svd(tensor_t, full_matrices=False)
        # 取前 target_rank 個奇異值
        U_reduced = U[:, :target_rank]  # (in_features, target_rank)
        S_reduced = S[:target_rank]  # (target_rank,)
        Vh_reduced = Vh[:target_rank, :]  # (target_rank, current_rank)
        # 重建: U @ diag(S) @ Vh，然後轉置回 (target_rank, in_features)
        aligned = (U_reduced @ torch.diag(S_reduced) @ Vh_reduced).t()
    else:
        # lora_B: (out_features, current_rank)
        # 直接進行 SVD
        U, S, Vh = torch.linalg.svd(tensor_fp32, full_matrices=False)
        # 取前 target_rank 個奇異值
        U_reduced = U[:, :target_rank]  # (out_features, target_rank)
        S_reduced = S[:target_rank]  # (target_rank,)
        Vh_reduced = Vh[:target_rank, :]  # (target_rank, current_rank)
        # 重建: U @ diag(S) @ Vh
        aligned = U_reduced @ torch.diag(S_reduced) @ Vh_reduced

    # 轉回原始數據類型
    return aligned.to(dtype)


def align_lora_rank(tensor, key, target_rank, target_shape, dtype, device, reduce_method='truncate'):
    """
    將 LoRA 張量的 rank 維度調整到目標值
    - lora_A.weight: rank 是第一個維度
    - lora_B.weight: rank 是最後一個維度

    reduce_method: 當目標 rank < 當前 rank 時的降維方法
        - 'truncate': 截斷，只取前 N 個維度
        - 'average': 平均聚合，將多個維度平均成 N 個維度
        - 'svd': 使用 SVD 進行低秩近似
    """
    rank_info = get_lora_rank_info(key, tensor.shape)
    if rank_info is None:
        # 不是 LoRA 鍵，使用通用對齊
        return align_tensor_shape(tensor, target_shape, dtype, device)

    rank_dim_index, current_rank = rank_info

    # 從 target_shape 中獲取目標 rank（應該已經被設置為指定的 dim）
    if rank_dim_index == 0:
        target_rank_from_shape = target_shape[0]
    else:
        target_rank_from_shape = target_shape[-1]

    # 確保 target_rank 與 target_shape 中的 rank 一致
    if target_rank != target_rank_from_shape:
        # 如果目標形狀中的 rank 與指定值不一致，使用指定值
        target_shape_list = list(target_shape)
        if rank_dim_index == 0:
            target_shape_list[0] = target_rank
        else:
            target_shape_list[-1] = target_rank
        target_shape = tuple(target_shape_list)

    if current_rank == target_rank and tensor.shape == target_shape:
        # rank 和形狀都已經相符
        return tensor

    if current_rank == target_rank:
        # rank 已經相符，但可能其他維度需要對齊
        return align_tensor_shape(tensor, target_shape, dtype, device)

    # 構建目標形狀（確保 rank 維度正確）
    target_shape_list = list(target_shape)
    if rank_dim_index == 0:
        target_shape_list[0] = target_rank
    else:
        target_shape_list[-1] = target_rank
    target_shape = tuple(target_shape_list)

    if current_rank > target_rank:
        # 降維：根據 reduce_method 選擇方法
        if reduce_method == 'truncate':
            # 截斷：只取前 target_rank 個維度
            if rank_dim_index == 0:
                aligned = tensor[:target_rank, ...]
            else:
                aligned = tensor[..., :target_rank]
        elif reduce_method == 'average':
            # 平均聚合
            aligned = reduce_rank_average(tensor, rank_dim_index, current_rank, target_rank, dtype, device)
        elif reduce_method == 'svd':
            # SVD 低秩近似
            aligned = reduce_rank_svd(tensor, rank_dim_index, current_rank, target_rank, dtype, device)
        else:
            raise ValueError(f"未知的降維方法: {reduce_method}")

        # 如果其他維度也需要對齊，使用通用對齊
        if aligned.shape != target_shape:
            aligned = align_tensor_shape(aligned, target_shape, dtype, device)
    else:
        # 擴展：用零填充
        aligned = torch.zeros(target_shape, dtype=dtype, device=device)
        if rank_dim_index == 0:
            # lora_A: 將原始數據放在前 current_rank 行
            aligned[:current_rank, ...] = tensor
        else:
            # lora_B: 將原始數據放在前 current_rank 列
            aligned[..., :current_rank] = tensor

    return aligned


# ============================================================================
# 參數解析和驗證
# ============================================================================

def parse_layer_weights(param_value, param_name, num_loras):
    """
    通用函數：解析 layer 權重參數（blocks, adaln, ff, attn 等）

    Args:
        param_value: 參數值列表（從 argparse）
        param_name: 參數名稱（用於錯誤消息）
        num_loras: LoRA 文件數量

    Returns:
        解析後的權重配置列表，每個元素是一個包含 30 個權重的列表
    """
    if param_value is None:
        return None

    if len(param_value) != num_loras:
        print(f"錯誤：--{param_name} 參數數量 ({len(param_value)}) 與 --lora 參數數量 ({num_loras}) 不相符")
        sys.exit(1)

    configs = []
    for i, weight_str in enumerate(param_value):
        try:
            weights = [float(x.strip()) for x in weight_str.split(',')]
        except ValueError as e:
            print(f"錯誤：無法解析第 {i+1} 個 LoRA 的 {param_name} 參數：{weight_str}")
            print(f"  原因：{e}")
            sys.exit(1)

        # 驗證權重數量必須為 30
        if len(weights) != 30:
            print(f"錯誤：第 {i+1} 個 LoRA 的 {param_name} 參數必須包含 30 個數字，但實際有 {len(weights)} 個")
            sys.exit(1)

        configs.append(weights)
        print(f"  解析 LoRA {i+1} {param_name} 權重：{weights[:5]}... (共 30 個)")

    return configs


def validate_args(args):
    """驗證命令行參數"""
    # 驗證參數數量相符
    if len(args.lora) != len(args.ratio):
        print(f"錯誤：--lora 參數數量 ({len(args.lora)}) 與 --ratio 參數數量 ({len(args.ratio)}) 不相符")
        sys.exit(1)

    if len(args.lora) < 1:
        print("錯誤：至少需要指定一個 LoRA 文件")
        sys.exit(1)

    # blocks 與 adaln, ff, attn, attn_to_k, attn_to_q, attn_to_v 不可混用
    incompatible_params = [args.adaln, args.ff, args.attn, args.attn_to_k, args.attn_to_q, args.attn_to_v]
    if args.blocks is not None and any(p is not None for p in incompatible_params):
        print("錯誤：--blocks 與 --adaln, --ff, --attn, --attn_to_k, --attn_to_q, --attn_to_v 不可混用")
        sys.exit(1)


def parse_all_layer_weights(args):
    """解析所有 layer 權重參數"""
    num_loras = len(args.lora)

    return {
        'blocks': parse_layer_weights(args.blocks, 'blocks', num_loras),
        'adaln': parse_layer_weights(args.adaln, 'adaln', num_loras),
        'ff': parse_layer_weights(args.ff, 'ff', num_loras),
        'attn': parse_layer_weights(args.attn, 'attn', num_loras),
        'attn_to_k': parse_layer_weights(args.attn_to_k, 'attn_to_k', num_loras),
        'attn_to_q': parse_layer_weights(args.attn_to_q, 'attn_to_q', num_loras),
        'attn_to_v': parse_layer_weights(args.attn_to_v, 'attn_to_v', num_loras),
        'attn_to_out': parse_layer_weights(args.attn_to_out, 'attn_to_out', num_loras),
    }


def prepare_lora_configs(args):
    """準備 LoRA 文件配置"""
    lora_configs = []
    for lora_path, weight in zip(args.lora, args.ratio):
        if not os.path.exists(lora_path):
            print(f"錯誤：LoRA 文件不存在：{lora_path}")
            sys.exit(1)
        lora_configs.append((lora_path, weight))
    return lora_configs


# ============================================================================
# 模型載入和驗證
# ============================================================================

def load_models(lora_configs):
    """載入所有 LoRA 模型"""
    print(f"載入 {len(lora_configs)} 個 LoRA 模型...")
    models = []
    for lora_path, weight in lora_configs:
        print(f"  載入：{lora_path} (權重: {weight})")
        models.append((load_file(lora_path), weight))
    return models


def validate_model_keys(models):
    """驗證所有模型的鍵是否一致"""
    first_model_keys = set(models[0][0].keys())

    for i, (model_sd, _) in enumerate(models[1:], 1):
        model_keys = set(model_sd.keys())
        if model_keys != first_model_keys:
            missing_keys = first_model_keys - model_keys
            extra_keys = model_keys - first_model_keys
            error_msg = f"錯誤：模型 {i+1} 的鍵與第一個模型不相符"
            if missing_keys:
                error_msg += f"\n  缺少的鍵：{missing_keys}"
            if extra_keys:
                error_msg += f"\n  多餘的鍵：{extra_keys}"
            print(error_msg)
            sys.exit(1)

    return first_model_keys


def detect_dtype(models, first_model_keys):
    """檢測第一個模型的數據類型"""
    first_key = next(iter(first_model_keys))
    target_dtype = models[0][0][first_key].dtype
    print(f"檢測到模型數據類型：{target_dtype}")
    return target_dtype


# ============================================================================
# 形狀檢查和準備
# ============================================================================

def compute_max_shapes(models, first_model_keys, args, layer_weight_configs):
    """
    檢查並確定每個鍵的最大形狀（處理不同 rank 的情況）

    Returns:
        tuple: (max_shapes, shape_info, layer_shapes_dict)
        - max_shapes: 每個鍵的最大形狀
        - shape_info: 每個鍵的所有形狀資訊（用於錯誤報告）
        - layer_shapes_dict: 按類型分組的形狀字典（adaln_shapes, ff_shapes 等）
    """
    print("檢查張量形狀...")
    max_shapes = {}
    shape_info = {}

    # 初始化各類型的形狀字典
    layer_shapes_dict = {
        'adaln': {},
        'ff': {},
        'attn': {},
        'attn_to_k': {},
        'attn_to_q': {},
        'attn_to_v': {},
        'attn_to_out': {},
    }

    for key in first_model_keys:
        shapes = []
        max_shape = None
        for i, (model_sd, _) in enumerate(models):
            tensor_shape = model_sd[key].shape
            shapes.append((i+1, tensor_shape))

            if max_shape is None:
                max_shape = tensor_shape
            else:
                # 檢查維度數量是否一致
                if len(tensor_shape) != len(max_shape):
                    print(f"錯誤：鍵 '{key}' 的張量維度數量不一致")
                    for idx, shape in shapes:
                        print(f"  模型 {idx}: {shape}")
                    sys.exit(1)
                # 逐個維度取最大值
                max_shape = tuple(max(max_shape[j], tensor_shape[j]) for j in range(len(max_shape)))

        # 如果指定了 --dim 且是 LoRA 鍵，調整目標形狀的 rank 維度
        if args.dim is not None and is_lora_key(key):
            rank_info = get_lora_rank_info(key, max_shape)
            if rank_info is not None:
                rank_dim_index, current_rank = rank_info
                if current_rank != args.dim:
                    # 調整目標形狀的 rank 維度
                    max_shape_list = list(max_shape)
                    if rank_dim_index == 0:
                        max_shape_list[0] = args.dim
                    else:
                        max_shape_list[-1] = args.dim
                    max_shape = tuple(max_shape_list)
                    print(f"  鍵 '{key}': 將 rank 維度設置為 {args.dim} (原最大: {current_rank})")

        max_shapes[key] = max_shape
        shape_info[key] = shapes

        # 根據配置收集各類型的形狀
        if layer_weight_configs['adaln'] is not None and 'adaLN_modulation' in key:
            layer_shapes_dict['adaln'][key] = max_shape
        if layer_weight_configs['ff'] is not None and 'feed_forward' in key:
            layer_shapes_dict['ff'][key] = max_shape
        if layer_weight_configs['attn'] is not None and 'attention' in key:
            layer_shapes_dict['attn'][key] = max_shape
        if layer_weight_configs['attn_to_k'] is not None and 'attention.to_k' in key:
            layer_shapes_dict['attn_to_k'][key] = max_shape
        if layer_weight_configs['attn_to_q'] is not None and 'attention.to_q' in key:
            layer_shapes_dict['attn_to_q'][key] = max_shape
        if layer_weight_configs['attn_to_v'] is not None and 'attention.to_v' in key:
            layer_shapes_dict['attn_to_v'][key] = max_shape
        if layer_weight_configs['attn_to_out'] is not None and 'attention.to_out' in key:
            layer_shapes_dict['attn_to_out'][key] = max_shape

        # 檢查是否有形狀不一致的情況
        all_same = all(s[1] == max_shape for s in shapes)
        if not all_same:
            print(f"  鍵 '{key}': 檢測到不同形狀，將對齊到最大形狀 {max_shape}")
            for idx, shape in shapes:
                if shape != max_shape:
                    print(f"    模型 {idx}: {shape} -> {max_shape}")

    return max_shapes, shape_info, layer_shapes_dict


# ============================================================================
# 模型合併
# ============================================================================

def calculate_actual_weight(key, layer_num, weight, i, layer_weight_configs, layer_shapes_dict, max_shapes):
    """
    計算實際權重，考慮 layer 權重配置

    Returns:
        tuple: (actual_weight, target_shape)
    """
    actual_weight = weight
    target_shape = max_shapes[key]

    # 如果有 blocks 配置且 key 屬於某個 layer，使用 blocks 權重
    if layer_weight_configs['blocks'] is not None and layer_num is not None:
        if 0 <= layer_num < 30:
            actual_weight = weight * layer_weight_configs['blocks'][i][layer_num]
    else:
        # 使用細粒度的 layer 權重配置
        if (any(cfg is not None for cfg in [
            layer_weight_configs['adaln'], layer_weight_configs['ff'],
            layer_weight_configs['attn'], layer_weight_configs['attn_to_k'],
            layer_weight_configs['attn_to_q'], layer_weight_configs['attn_to_v'],
            layer_weight_configs['attn_to_out']
        ]) and layer_num is not None):

            if key in layer_shapes_dict['adaln']:
                target_shape = layer_shapes_dict['adaln'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['adaln'][i][layer_num]
            elif key in layer_shapes_dict['ff']:
                target_shape = layer_shapes_dict['ff'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['ff'][i][layer_num]
            elif key in layer_shapes_dict['attn']:
                target_shape = layer_shapes_dict['attn'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['attn'][i][layer_num]
            elif key in layer_shapes_dict['attn_to_k']:
                target_shape = layer_shapes_dict['attn_to_k'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['attn_to_k'][i][layer_num]
            elif key in layer_shapes_dict['attn_to_q']:
                target_shape = layer_shapes_dict['attn_to_q'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['attn_to_q'][i][layer_num]
            elif key in layer_shapes_dict['attn_to_v']:
                target_shape = layer_shapes_dict['attn_to_v'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['attn_to_v'][i][layer_num]
            elif key in layer_shapes_dict['attn_to_out']:
                target_shape = layer_shapes_dict['attn_to_out'][key]
                if 0 <= layer_num < 30:
                    actual_weight = weight * layer_weight_configs['attn_to_out'][i][layer_num]

    return actual_weight, target_shape


def merge_models(models, first_model_keys, max_shapes, layer_weight_configs,
                 layer_shapes_dict, target_dtype, args):
    """合併所有模型"""
    print("合併模型...")
    if args.dim is not None:
        print(f"  使用指定的 rank 維度：{args.dim}")
        print(f"  降維方法：{args.dim_reduce_method}")

    # 印出使用的配置資訊
    config_names = {
        'blocks': 'block 權重配置',
        'adaln': 'adaLN_modulation 權重配置',
        'ff': 'feed_forward 權重配置',
        'attn': 'attention 權重配置',
        'attn_to_k': 'attention.to_k 權重配置',
        'attn_to_q': 'attention.to_q 權重配置',
        'attn_to_v': 'attention.to_v 權重配置',
        'attn_to_out': 'attention.to_out 權重配置',
    }
    for config_key, config_name in config_names.items():
        if layer_weight_configs[config_key] is not None:
            print(f"  使用 {config_name}")

    new_model_sd = {}
    for key in first_model_keys:
        # 提取 layer 編號（如果適用）
        layer_num = extract_layer_number(key)
        target_shape = max_shapes[key]

        # 加權求和，保持原始數據類型
        merged_tensor = None
        for i, (model_sd, weight) in enumerate(models):
            tensor = model_sd[key]

            # 計算實際權重
            actual_weight, target_shape = calculate_actual_weight(
                key, layer_num, weight, i, layer_weight_configs,
                layer_shapes_dict, max_shapes
            )

            # 將權重轉換為與張量相同的數據類型，避免類型提升
            weight_tensor = torch.tensor(actual_weight, dtype=target_dtype, device=tensor.device)

            # 對齊張量形狀
            # 如果指定了 --dim 且是 LoRA 鍵，使用專門的 rank 對齊函數
            if args.dim is not None and is_lora_key(key):
                aligned_tensor = align_lora_rank(
                    tensor, key, args.dim, target_shape, target_dtype,
                    tensor.device, args.dim_reduce_method
                )
            else:
                aligned_tensor = align_tensor_shape(tensor, target_shape, target_dtype, tensor.device)

            if merged_tensor is None:
                merged_tensor = aligned_tensor * weight_tensor
            else:
                merged_tensor = merged_tensor + aligned_tensor * weight_tensor

        # 確保結果保持原始數據類型（防止運算過程中的類型提升）
        merged_tensor = merged_tensor.to(target_dtype)
        new_model_sd[key] = merged_tensor

    return new_model_sd


# ============================================================================
# 數據類型轉換和儲存
# ============================================================================

def convert_dtype(new_model_sd, args, target_dtype):
    """根據參數轉換數據類型"""
    if args.dtype != 'auto':
        print(f"轉換模型數據類型為：{args.dtype}")
        if args.dtype == 'bf16':
            output_dtype = torch.bfloat16
        elif args.dtype == 'fp16':
            output_dtype = torch.float16
        else:
            output_dtype = target_dtype

        for key in new_model_sd:
            # 只轉換浮點數類型的張量
            if new_model_sd[key].dtype in [torch.float32, torch.float16, torch.bfloat16]:
                new_model_sd[key] = new_model_sd[key].to(output_dtype)
        print(f"轉換完成，輸出數據類型：{output_dtype}")


def save_model(new_model_sd, output_path):
    """儲存模型"""
    print(f"儲存到：{output_path}")
    save_file(new_model_sd, output_path)
    print("完成！")


# ============================================================================
# 主程式
# ============================================================================

def create_parser():
    """建立命令行參數解析器"""
    parser = argparse.ArgumentParser(description="合併多個使用 AI-Toolkit 訓練的 Z Image Turbo LoRA 模型")
    parser.add_argument(
        '--lora',
        nargs='+',
        required=True,
        help='LoRA 文件路徑，可指定多個文件，例如：--lora file1.safetensors file2.safetensors file3.safetensors'
    )
    parser.add_argument(
        '--ratio',
        nargs='+',
        required=True,
        type=float,
        help='LoRA 權重比例，可指定多個權重，例如：--ratio 0.1 0.3 0.5。必須與 --lora 參數數量相同'
    )
    parser.add_argument(
        '--output',
        type=str,
        required=True,
        help='輸出文件路徑'
    )
    parser.add_argument(
        '--dtype',
        type=str,
        choices=['auto', 'bf16', 'fp16'],
        default='auto',
        help='輸出模型的數據類型：auto（保持原樣）、bf16（bfloat16）、fp16（float16）。默認為 auto'
    )
    parser.add_argument(
        '--blocks',
        nargs='+',
        default=None,
        help='每個 LoRA 的 block 權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 blocks 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--adaln',
        nargs='+',
        default=None,
        help='特別指定關於 adaLN_modulation 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 adaLN_modulation 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--ff',
        nargs='+',
        default=None,
        help='特別指定關於 feed_forward 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 feed_forward 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--attn',
        nargs='+',
        default=None,
        help='特別指定關於 attention 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 attention 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--attn_to_k',
        nargs='+',
        default=None,
        help='特別指定關於 attention.to_k 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 attention.to_k 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--attn_to_q',
        nargs='+',
        default=None,
        help='特別指定關於 attention.to_q 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 attention.to_q 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--attn_to_v',
        nargs='+',
        default=None,
        help='特別指定關於 attention.to_v 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 attention.to_v 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--attn_to_out',
        nargs='+',
        default=None,
        help='特別指定關於 attention.to_out 的權重，格式為 "0,1,0,1..." 共 30 個數字（對應 layers.0 到 layers.29）。可指定多個，必須與 --lora 參數數量相同。如果未指定，則所有 attention.to_out 使用相同的 ratio 權重'
    )
    parser.add_argument(
        '--dim',
        type=int,
        default=None,
        help='指定 LoRA 的 rank 維度。如果指定，所有 lora_A.weight 和 lora_B.weight 的 rank 維度將轉換為此值。lora_A.weight 的 rank 是第一個維度，lora_B.weight 的 rank 是最後一個維度'
    )
    parser.add_argument(
        '--dim_reduce_method',
        type=str,
        choices=['truncate', 'average', 'svd'],
        default='truncate',
        help='當目標 rank 小於當前 rank 時的降維方法：truncate（截斷，只取前 N 個維度）、average（平均聚合，將多個維度平均成 N 個維度）、svd（使用 SVD 進行低秩近似）。默認為 truncate'
    )

    return parser


if __name__ == "__main__":
    # 解析命令行參數
    parser = create_parser()
    args = parser.parse_args()

    # 驗證參數
    validate_args(args)

    # 解析所有 layer 權重參數
    layer_weight_configs = parse_all_layer_weights(args)

    # 準備 LoRA 文件配置
    lora_configs = prepare_lora_configs(args)

    # 載入模型
    models = load_models(lora_configs)

    # 驗證模型鍵一致性
    first_model_keys = validate_model_keys(models)

    # 檢測數據類型
    target_dtype = detect_dtype(models, first_model_keys)

    # 計算最大形狀
    max_shapes, shape_info, layer_shapes_dict = compute_max_shapes(
        models, first_model_keys, args, layer_weight_configs
    )

    # 合併模型
    new_model_sd = merge_models(
        models, first_model_keys, max_shapes, layer_weight_configs,
        layer_shapes_dict, target_dtype, args
    )

    # 轉換數據類型
    convert_dtype(new_model_sd, args, target_dtype)

    # 儲存模型
    save_model(new_model_sd, args.output)

modify_zlora_en.py

import argparse
import re
import torch
from safetensors.torch import load_file, save_file
import sys
import os


# ============================================================================
# Helper Functions: Key and Layer Processing
# ============================================================================

def extract_layer_number(key):
    """Extracts the layer number from a key, e.g., 'layers.0.xxx' -> 0, 'layers.29.xxx' -> 29"""
    match = re.search(r'layers\.(\d+)', key)
    if match:
        return int(match.group(1))
    return None


def is_lora_key(key):
    """Checks if the key is a LoRA weight (lora_A.weight or lora_B.weight)"""
    return 'lora_A.weight' in key or 'lora_B.weight' in key


def get_lora_rank_info(key, shape):
    """
    Retrieves the rank dimension information for a LoRA key.
    Returns: (rank_dim_index, current_rank) or None if not a LoRA key.
    """
    if 'lora_A.weight' in key:
        # For lora_A.weight, rank is the first dimension
        return (0, shape[0])
    elif 'lora_B.weight' in key:
        # For lora_B.weight, rank is the last dimension
        return (-1, shape[-1])
    return None


# ============================================================================
# Helper Functions: Tensor Alignment and Rank Adjustment
# ============================================================================

def align_tensor_shape(tensor, target_shape, dtype, device):
    """Adjusts a tensor to the target shape, padding missing parts with zeros."""
    if tensor.shape == target_shape:
        return tensor

    # Check if number of dimensions match
    if len(tensor.shape) != len(target_shape):
        raise ValueError(
            f"Unable to align tensor: dimension mismatch "
            f"({len(tensor.shape)}D vs {len(target_shape)}D). "
            f"Shapes: {tensor.shape} vs {target_shape}"
        )

    # Create a zero tensor with the target shape
    aligned = torch.zeros(target_shape, dtype=dtype, device=device)

    # Build slices to copy data to the top-left of the aligned tensor
    # For each dimension, take the minimum of original and target size
    slices = tuple(slice(0, min(tensor.shape[i], target_shape[i])) for i in range(len(target_shape)))

    # Copy data from the original tensor
    aligned[slices] = tensor[slices]

    return aligned


def reduce_rank_average(tensor, rank_dim_index, current_rank, target_rank, dtype, device):
    """
    Reduces rank from current_rank to target_rank using an average aggregation method.
    Groups dimensions and calculates the mean.
    """
    if rank_dim_index == 0:
        # lora_A: (current_rank, ...) -> (target_rank, ...)
        group_size = current_rank / target_rank
        aligned = torch.zeros((target_rank,) + tensor.shape[1:], dtype=dtype, device=device)

        for i in range(target_rank):
            start_idx = int(i * group_size)
            end_idx = int((i + 1) * group_size) if i < target_rank - 1 else current_rank
            # Calculate mean for this group
            aligned[i, ...] = tensor[start_idx:end_idx, ...].mean(dim=0)
    else:
        # lora_B: (..., current_rank) -> (..., target_rank)
        group_size = current_rank / target_rank
        aligned = torch.zeros(tensor.shape[:-1] + (target_rank,), dtype=dtype, device=device)

        for i in range(target_rank):
            start_idx = int(i * group_size)
            end_idx = int((i + 1) * group_size) if i < target_rank - 1 else current_rank
            # Calculate mean for this group
            aligned[..., i] = tensor[..., start_idx:end_idx].mean(dim=-1)

    return aligned


def reduce_rank_svd(tensor, rank_dim_index, current_rank, target_rank, dtype, device):
    """
    Reduces rank using SVD (Singular Value Decomposition) for low-rank approximation.
    Note: SVD requires float32 precision for stability.
    """
    # SVD requires float32, convert first
    tensor_fp32 = tensor.float()

    if rank_dim_index == 0:
        # lora_A: (current_rank, in_features)
        # Transpose to (in_features, current_rank) for SVD
        tensor_t = tensor_fp32.t()  # (in_features, current_rank)
        U, S, Vh = torch.linalg.svd(tensor_t, full_matrices=False)
        # Keep top target_rank singular values
        U_reduced = U[:, :target_rank]  # (in_features, target_rank)
        S_reduced = S[:target_rank]  # (target_rank,)
        Vh_reduced = Vh[:target_rank, :]  # (target_rank, current_rank)
        # Reconstruct and transpose back to (target_rank, in_features)
        aligned = (U_reduced @ torch.diag(S_reduced) @ Vh_reduced).t()
    else:
        # lora_B: (out_features, current_rank)
        U, S, Vh = torch.linalg.svd(tensor_fp32, full_matrices=False)
        # Keep top target_rank singular values
        U_reduced = U[:, :target_rank]  # (out_features, target_rank)
        S_reduced = S[:target_rank]  # (target_rank,)
        Vh_reduced = Vh[:target_rank, :]  # (target_rank, current_rank)
        # Reconstruct
        aligned = U_reduced @ torch.diag(S_reduced) @ Vh_reduced

    # Convert back to original dtype
    return aligned.to(dtype)


def align_lora_rank(tensor, key, target_rank, target_shape, dtype, device, reduce_method='truncate'):
    """
    Adjusts the rank dimension of a LoRA tensor to a target value.
    - lora_A.weight: rank is the first dimension.
    - lora_B.weight: rank is the last dimension.

    reduce_method: Strategy when target rank < current rank:
        - 'truncate': Simply slice the first N dimensions.
        - 'average': Group and average dimensions.
        - 'svd': Low-rank approximation using SVD.
    """
    rank_info = get_lora_rank_info(key, tensor.shape)
    if rank_info is None:
        # Not a LoRA key, use general alignment
        return align_tensor_shape(tensor, target_shape, dtype, device)

    rank_dim_index, current_rank = rank_info

    # Determine target rank from shape (should match the specified dimension)
    if rank_dim_index == 0:
        target_rank_from_shape = target_shape[0]
    else:
        target_rank_from_shape = target_shape[-1]

    # Sync target_rank with target_shape
    if target_rank != target_rank_from_shape:
        target_shape_list = list(target_shape)
        if rank_dim_index == 0:
            target_shape_list[0] = target_rank
        else:
            target_shape_list[-1] = target_rank
        target_shape = tuple(target_shape_list)

    if current_rank == target_rank and tensor.shape == target_shape:
        return tensor

    if current_rank == target_rank:
        # Rank matches, but other dimensions may need alignment
        return align_tensor_shape(tensor, target_shape, dtype, device)

    # Re-verify target shape consistency
    target_shape_list = list(target_shape)
    if rank_dim_index == 0:
        target_shape_list[0] = target_rank
    else:
        target_shape_list[-1] = target_rank
    target_shape = tuple(target_shape_list)

    if current_rank > target_rank:
        # Rank Reduction
        if reduce_method == 'truncate':
            if rank_dim_index == 0:
                aligned = tensor[:target_rank, ...]
            else:
                aligned = tensor[..., :target_rank]
        elif reduce_method == 'average':
            aligned = reduce_rank_average(tensor, rank_dim_index, current_rank, target_rank, dtype, device)
        elif reduce_method == 'svd':
            aligned = reduce_rank_svd(tensor, rank_dim_index, current_rank, target_rank, dtype, device)
        else:
            raise ValueError(f"Unknown reduction method: {reduce_method}")

        # Align other dimensions if necessary
        if aligned.shape != target_shape:
            aligned = align_tensor_shape(aligned, target_shape, dtype, device)
    else:
        # Expansion: Pad with zeros
        aligned = torch.zeros(target_shape, dtype=dtype, device=device)
        if rank_dim_index == 0:
            # lora_A: Place original data in the top rows
            aligned[:current_rank, ...] = tensor
        else:
            # lora_B: Place original data in the first columns
            aligned[..., :current_rank] = tensor

    return aligned


# ============================================================================
# Argument Parsing and Validation
# ============================================================================

def parse_layer_weights(param_value, param_name, num_loras):
    """
    Generic parser for layer weight parameters (blocks, adaln, ff, attn, etc.)
    """
    if param_value is None:
        return None

    if len(param_value) != num_loras:
        print(f"Error: Number of --{param_name} arguments ({len(param_value)}) does not match --lora count ({num_loras})")
        sys.exit(1)

    configs = []
    for i, weight_str in enumerate(param_value):
        try:
            weights = [float(x.strip()) for x in weight_str.split(',')]
        except ValueError as e:
            print(f"Error: Could not parse {param_name} for LoRA {i+1}: {weight_str}")
            print(f"  Reason: {e}")
            sys.exit(1)

        # Validate weight count (must be 30 for Z-Image-Turbo/AI-Toolkit standard)
        if len(weights) != 30:
            print(f"Error: {param_name} for LoRA {i+1} must contain exactly 30 numbers (found {len(weights)})")
            sys.exit(1)

        configs.append(weights)
        print(f"  Parsed LoRA {i+1} {param_name} weights: {weights[:5]}... (Total 30)")

    return configs


def validate_args(args):
    """Validates command line arguments."""
    if len(args.lora) != len(args.ratio):
        print(f"Error: Number of --lora files ({len(args.lora)}) does not match --ratio count ({len(args.ratio)})")
        sys.exit(1)

    if len(args.lora) < 1:
        print("Error: At least one LoRA file must be specified")
        sys.exit(1)

    # Prevent conflict between global 'blocks' and specific layer parameters
    incompatible_params = [args.adaln, args.ff, args.attn, args.attn_to_k, args.attn_to_q, args.attn_to_v]
    if args.blocks is not None and any(p is not None for p in incompatible_params):
        print("Error: --blocks cannot be used simultaneously with specific layer parameters (--adaln, --ff, --attn, etc.)")
        sys.exit(1)


def parse_all_layer_weights(args):
    """Parses all detailed layer weight parameters."""
    num_loras = len(args.lora)

    return {
        'blocks': parse_layer_weights(args.blocks, 'blocks', num_loras),
        'adaln': parse_layer_weights(args.adaln, 'adaln', num_loras),
        'ff': parse_layer_weights(args.ff, 'ff', num_loras),
        'attn': parse_layer_weights(args.attn, 'attn', num_loras),
        'attn_to_k': parse_layer_weights(args.attn_to_k, 'attn_to_k', num_loras),
        'attn_to_q': parse_layer_weights(args.attn_to_q, 'attn_to_q', num_loras),
        'attn_to_v': parse_layer_weights(args.attn_to_v, 'attn_to_v', num_loras),
        'attn_to_out': parse_layer_weights(args.attn_to_out, 'attn_to_out', num_loras),
    }


def prepare_lora_configs(args):
    """Prepares list of LoRA paths and their associated global ratios."""
    lora_configs = []
    for lora_path, weight in zip(args.lora, args.ratio):
        if not os.path.exists(lora_path):
            print(f"Error: LoRA file not found: {lora_path}")
            sys.exit(1)
        lora_configs.append((lora_path, weight))
    return lora_configs


# ============================================================================
# Model Loading and Verification
# ============================================================================

def load_models(lora_configs):
    """Loads all LoRA models into memory."""
    print(f"Loading {len(lora_configs)} LoRA model(s)...")
    models = []
    for lora_path, weight in lora_configs:
        print(f"  Loading: {lora_path} (Global Ratio: {weight})")
        models.append((load_file(lora_path), weight))
    return models


def validate_model_keys(models):
    """Ensures all models share the same state dict keys."""
    first_model_keys = set(models[0][0].keys())

    for i, (model_sd, _) in enumerate(models[1:], 1):
        model_keys = set(model_sd.keys())
        if model_keys != first_model_keys:
            missing_keys = first_model_keys - model_keys
            extra_keys = model_keys - first_model_keys
            error_msg = f"Error: Keys in Model {i+1} do not match the first model"
            if missing_keys:
                error_msg += f"\n  Missing keys: {missing_keys}"
            if extra_keys:
                error_msg += f"\n  Extra keys: {extra_keys}"
            print(error_msg)
            sys.exit(1)

    return first_model_keys


def detect_dtype(models, first_model_keys):
    """Detects the data type of the first model."""
    first_key = next(iter(first_model_keys))
    target_dtype = models[0][0][first_key].dtype
    print(f"Detected model data type: {target_dtype}")
    return target_dtype


# ============================================================================
# Shape Checking and Preparation
# ============================================================================

def compute_max_shapes(models, first_model_keys, args, layer_weight_configs):
    """
    Analyzes and determines the maximum shape for each key (handling mixed ranks).
    """
    print("Inspecting tensor shapes...")
    max_shapes = {}
    shape_info = {}

    # Group shapes by type
    layer_shapes_dict = {
        'adaln': {},
        'ff': {},
        'attn': {},
        'attn_to_k': {},
        'attn_to_q': {},
        'attn_to_v': {},
        'attn_to_out': {},
    }

    for key in first_model_keys:
        shapes = []
        max_shape = None
        for i, (model_sd, _) in enumerate(models):
            tensor_shape = model_sd[key].shape
            shapes.append((i+1, tensor_shape))

            if max_shape is None:
                max_shape = tensor_shape
            else:
                if len(tensor_shape) != len(max_shape):
                    print(f"Error: Dimension mismatch for key '{key}'")
                    for idx, shape in shapes:
                        print(f"  Model {idx}: {shape}")
                    sys.exit(1)
                # Take the maximum size for each dimension
                max_shape = tuple(max(max_shape[j], tensor_shape[j]) for j in range(len(max_shape)))

        # Adjust rank dimension if --dim is specified
        if args.dim is not None and is_lora_key(key):
            rank_info = get_lora_rank_info(key, max_shape)
            if rank_info is not None:
                rank_dim_index, current_rank = rank_info
                if current_rank != args.dim:
                    max_shape_list = list(max_shape)
                    if rank_dim_index == 0:
                        max_shape_list[0] = args.dim
                    else:
                        max_shape_list[-1] = args.dim
                    max_shape = tuple(max_shape_list)
                    print(f"  Key '{key}': Setting rank to {args.dim} (Original max: {current_rank})")

        max_shapes[key] = max_shape
        shape_info[key] = shapes

        # Categorize keys for fine-grained weighting
        if layer_weight_configs['adaln'] is not None and 'adaLN_modulation' in key:
            layer_shapes_dict['adaln'][key] = max_shape
        if layer_weight_configs['ff'] is not None and 'feed_forward' in key:
            layer_shapes_dict['ff'][key] = max_shape
        if layer_weight_configs['attn'] is not None and 'attention' in key:
            layer_shapes_dict['attn'][key] = max_shape
        if layer_weight_configs['attn_to_k'] is not None and 'attention.to_k' in key:
            layer_shapes_dict['attn_to_k'][key] = max_shape
        if layer_weight_configs['attn_to_q'] is not None and 'attention.to_q' in key:
            layer_shapes_dict['attn_to_q'][key] = max_shape
        if layer_weight_configs['attn_to_v'] is not None and 'attention.to_v' in key:
            layer_shapes_dict['attn_to_v'][key] = max_shape
        if layer_weight_configs['attn_to_out'] is not None and 'attention.to_out' in key:
            layer_shapes_dict['attn_to_out'][key] = max_shape

        # Log shape discrepancies
        all_same = all(s[1] == max_shape for s in shapes)
        if not all_same:
            print(f"  Key '{key}': Mixed shapes detected. Aligning to maximum: {max_shape}")
            for idx, shape in shapes:
                if shape != max_shape:
                    print(f"    Model {idx}: {shape} -> {max_shape}")

    return max_shapes, shape_info, layer_shapes_dict


# ============================================================================
# Model Merging Logic
# ============================================================================

def calculate_actual_weight(key, layer_num, weight, i, layer_weight_configs, layer_shapes_dict, max_shapes):
    """
    Calculates the effective weight based on layer-specific configurations.
    """
    actual_weight = weight
    target_shape = max_shapes[key]

    # Handle global block config
    if layer_weight_configs['blocks'] is not None and layer_num is not None:
        if 0 <= layer_num < 30:
            actual_weight = weight * layer_weight_configs['blocks'][i][layer_num]
    else:
        # Handle fine-grained layer configs
        if (any(cfg is not None for cfg in [
            layer_weight_configs['adaln'], layer_weight_configs['ff'],
            layer_weight_configs['attn'], layer_weight_configs['attn_to_k'],
            layer_weight_configs['attn_to_q'], layer_weight_configs['attn_to_v'],
            layer_weight_configs['attn_to_out']
        ]) and layer_num is not None):

            config_mapping = [
                ('adaln', 'adaln'), ('ff', 'ff'), ('attn', 'attn'),
                ('attn_to_k', 'attn_to_k'), ('attn_to_q', 'attn_to_q'),
                ('attn_to_v', 'attn_to_v'), ('attn_to_out', 'attn_to_out')
            ]

            for config_key, dict_key in config_mapping:
                if key in layer_shapes_dict[dict_key]:
                    target_shape = layer_shapes_dict[dict_key][key]
                    if 0 <= layer_num < 30:
                        actual_weight = weight * layer_weight_configs[config_key][i][layer_num]
                    break

    return actual_weight, target_shape


def merge_models(models, first_model_keys, max_shapes, layer_weight_configs,
                 layer_shapes_dict, target_dtype, args):
    """Performs the weighted sum and merging of all models."""
    print("Merging models...")
    if args.dim is not None:
        print(f"  Using target rank dimension: {args.dim}")
        print(f"  Reduction method: {args.dim_reduce_method}")

    # Log active configurations
    config_names = {
        'blocks': 'Block weight config',
        'adaln': 'adaLN_modulation weight config',
        'ff': 'feed_forward weight config',
        'attn': 'attention weight config',
        'attn_to_k': 'attention.to_k weight config',
        'attn_to_q': 'attention.to_q weight config',
        'attn_to_v': 'attention.to_v weight config',
        'attn_to_out': 'attention.to_out weight config',
    }
    for config_key, config_name in config_names.items():
        if layer_weight_configs[config_key] is not None:
            print(f"  Active: {config_name}")

    new_model_sd = {}
    for key in first_model_keys:
        layer_num = extract_layer_number(key)
        target_shape = max_shapes[key]

        merged_tensor = None
        for i, (model_sd, weight) in enumerate(models):
            tensor = model_sd[key]

            actual_weight, target_shape = calculate_actual_weight(
                key, layer_num, weight, i, layer_weight_configs,
                layer_shapes_dict, max_shapes
            )

            # Use tensor for weight to prevent automatic type promotion
            weight_tensor = torch.tensor(actual_weight, dtype=target_dtype, device=tensor.device)

            # Rank and Shape Alignment
            if args.dim is not None and is_lora_key(key):
                aligned_tensor = align_lora_rank(
                    tensor, key, args.dim, target_shape, target_dtype,
                    tensor.device, args.dim_reduce_method
                )
            else:
                aligned_tensor = align_tensor_shape(tensor, target_shape, target_dtype, tensor.device)

            if merged_tensor is None:
                merged_tensor = aligned_tensor * weight_tensor
            else:
                merged_tensor = merged_tensor + aligned_tensor * weight_tensor

        # Ensure final tensor maintains target precision
        merged_tensor = merged_tensor.to(target_dtype)
        new_model_sd[key] = merged_tensor

    return new_model_sd


# ============================================================================
# Dtype Conversion and Output
# ============================================================================

def convert_dtype(new_model_sd, args, target_dtype):
    """Converts the final model to the specified data type."""
    if args.dtype != 'auto':
        print(f"Converting model data type to: {args.dtype}")
        if args.dtype == 'bf16':
            output_dtype = torch.bfloat16
        elif args.dtype == 'fp16':
            output_dtype = torch.float16
        else:
            output_dtype = target_dtype

        for key in new_model_sd:
            if new_model_sd[key].dtype in [torch.float32, torch.float16, torch.bfloat16]:
                new_model_sd[key] = new_model_sd[key].to(output_dtype)
        print(f"Conversion complete. Output dtype: {output_dtype}")


def save_model(new_model_sd, output_path):
    """Saves the state dict to a safetensors file."""
    print(f"Saving to: {output_path}")
    save_file(new_model_sd, output_path)
    print("Done!")


# ============================================================================
# CLI Setup
# ============================================================================

def create_parser():
    """Defines the command-line interface."""
    parser = argparse.ArgumentParser(description="Merge multiple Z-Image-Turbo LoRA models trained with AI-Toolkit")
    parser.add_argument(
        '--lora',
        nargs='+',
        required=True,
        help='Paths to LoRA files. Multiple files can be specified, e.g., --lora f1.safetensors f2.safetensors'
    )
    parser.add_argument(
        '--ratio',
        nargs='+',
        required=True,
        type=float,
        help='Global ratios for each LoRA. Must match the number of --lora files.'
    )
    parser.add_argument(
        '--output',
        type=str,
        required=True,
        help='Path for the output LoRA file'
    )
    parser.add_argument(
        '--dtype',
        type=str,
        choices=['auto', 'bf16', 'fp16'],
        default='auto',
        help='Output data type: auto (keep original), bf16, or fp16. Default: auto'
    )
    parser.add_argument(
        '--blocks',
        nargs='+',
        default=None,
        help='Block weights for each LoRA. Format: "0,1,0.5..." (30 values for layers 0-29).'
    )
    parser.add_argument(
        '--adaln',
        nargs='+',
        default=None,
        help='Specific weights for adaLN_modulation (Global Ambiance). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--ff',
        nargs='+',
        default=None,
        help='Specific weights for feed_forward (Knowledge/Details). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--attn',
        nargs='+',
        default=None,
        help='Specific weights for general attention layers. Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--attn_to_k',
        nargs='+',
        default=None,
        help='Specific weights for attention.to_k (Layout/Structure). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--attn_to_q',
        nargs='+',
        default=None,
        help='Specific weights for attention.to_q (Query). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--attn_to_v',
        nargs='+',
        default=None,
        help='Specific weights for attention.to_v (Features/Content). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--attn_to_out',
        nargs='+',
        default=None,
        help='Specific weights for attention.to_out (Integration). Format: 30 comma-separated values.'
    )
    parser.add_argument(
        '--dim',
        type=int,
        default=None,
        help='Target rank dimension. If set, all LoRA weights will be converted to this rank.'
    )
    parser.add_argument(
        '--dim_reduce_method',
        type=str,
        choices=['truncate', 'average', 'svd'],
        default='truncate',
        help='Strategy for rank reduction: truncate, average, or svd. Default: truncate'
    )

    return parser


if __name__ == "__main__":
    parser = create_parser()
    args = parser.parse_args()

    validate_args(args)
    layer_weight_configs = parse_all_layer_weights(args)
    lora_configs = prepare_lora_configs(args)

    models = load_models(lora_configs)
    first_model_keys = validate_model_keys(models)
    target_dtype = detect_dtype(models, first_model_keys)

    max_shapes, shape_info, layer_shapes_dict = compute_max_shapes(
        models, first_model_keys, args, layer_weight_configs
    )

    new_model_sd = merge_models(
        models, first_model_keys, max_shapes, layer_weight_configs,
        layer_shapes_dict, target_dtype, args
    )

    convert_dtype(new_model_sd, args, target_dtype)
    save_model(new_model_sd, args.output)