#!/usr/bin/env python3 import argparse import json import os from dataclasses import dataclass from typing import Any, Callable, Dict, List, Tuple, Optional import torch from safetensors import safe_open from safetensors.torch import save_file from tqdm import tqdm ProgressCallback = Callable[[str, int, int, Optional[str]], None] def _progress(cb: Optional[ProgressCallback], stage: str, current: int, total: int, message: Optional[str] = None): if cb is None: return try: cb(stage, int(current), int(total), message) except Exception: pass @dataclass class LoraPair: down: torch.Tensor up: torch.Tensor alpha: float is_conv: bool in_hw: Tuple[int, int] def parse_weighted_paths(items: List[str]) -> List[Tuple[str, float]]: result: List[Tuple[str, float]] = [] for item in items: if "@" in item: path, weight_str = item.rsplit("@", 1) weight = float(weight_str) result.append((path, weight)) else: result.append((item, 1.0)) return result def get_device(force_cpu: bool) -> torch.device: if force_cpu: return torch.device("cpu") if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") def get_dtype(dtype_name: str) -> torch.dtype: name = dtype_name.lower() if name in ("fp16", "float16", "half"): return torch.float16 if name in ("fp32", "float32"): return torch.float32 if name in ("bf16", "bfloat16"): return torch.bfloat16 raise ValueError(f"Unsupported dtype: {dtype_name}") def get_compute_dtype(dtype_name: str) -> Optional[torch.dtype]: """ Compute dtype is for internal merge math alignment. If None => auto behavior. """ name = dtype_name.lower() if name in ("auto", "none"): return None return get_dtype(name) def is_lora_key(key: str) -> bool: return ( key.endswith("lora_down.weight") or key.endswith("lora_up.weight") or key.endswith("lora_down") or key.endswith("lora_up") or key.endswith("lora_A.weight") or key.endswith("lora_B.weight") or key.endswith("lora_A") or key.endswith("lora_B") ) def get_module_prefix(key: str) -> str: if ".lora_down" in key: return key.split(".lora_down", 1)[0] if ".lora_up" in key: return key.split(".lora_up", 1)[0] if ".lora_A" in key: return key.split(".lora_A", 1)[0] if ".lora_B" in key: return key.split(".lora_B", 1)[0] parts = key.split(".") return ".".join(parts[:-2]) def get_alpha_key(prefix: str) -> str: return f"{prefix}.alpha" def _get_save_key_style(metadata_ref: Dict[str, str]) -> str: style = (metadata_ref or {}).get("_was_lora_key_style", "") style_norm = str(style).strip().lower() if style_norm in ("ab", "a/b", "lora_a_b"): return "ab" return "downup" def _down_key(prefix: str, key_style: str) -> str: if key_style == "ab": return f"{prefix}.lora_A.weight" return f"{prefix}.lora_down.weight" def _up_key(prefix: str, key_style: str) -> str: if key_style == "ab": return f"{prefix}.lora_B.weight" return f"{prefix}.lora_up.weight" def module_is_included( prefix: str, include_patterns: List[str], exclude_patterns: List[str], ) -> bool: if include_patterns: if not any(pat in prefix for pat in include_patterns): return False if exclude_patterns: if any(pat in prefix for pat in exclude_patterns): return False return True def load_lora_pairs( path: str, device: torch.device, progress_cb: Optional[ProgressCallback] = None, ) -> Tuple[Dict[str, LoraPair], Dict[str, str]]: pairs: Dict[str, LoraPair] = {} metadata: Dict[str, str] = {} with safe_open(path, framework="pt", device=str(device)) as f: keys = list(f.keys()) try: raw_meta = f.metadata() metadata = dict(raw_meta) if isinstance(raw_meta, dict) else {} except Exception: metadata = {} detected_key_style: Optional[str] = None for k in keys: if k.endswith("lora_A.weight") or k.endswith("lora_B.weight") or k.endswith("lora_A") or k.endswith("lora_B"): detected_key_style = "ab" break if detected_key_style is None: for k in keys: if k.endswith("lora_down.weight") or k.endswith("lora_up.weight") or k.endswith("lora_down") or k.endswith("lora_up"): detected_key_style = "downup" break if detected_key_style is not None: metadata["_was_lora_key_style"] = detected_key_style downs: Dict[str, torch.Tensor] = {} ups: Dict[str, torch.Tensor] = {} alphas: Dict[str, float] = {} total_keys = len(keys) for i, key in enumerate(keys, start=1): _progress(progress_cb, "load.keys", i, total_keys, key) if is_lora_key(key): tensor = f.get_tensor(key).to(device) prefix = get_module_prefix(key) if "lora_down" in key or "lora_A" in key: downs[prefix] = tensor elif "lora_up" in key or "lora_B" in key: ups[prefix] = tensor elif key.endswith(".alpha"): tensor = f.get_tensor(key) value = float(tensor.item()) try: prefix = get_module_prefix(key.replace(".alpha", ".lora_up")) alphas[prefix] = value except Exception: try: prefix = get_module_prefix(key.replace(".alpha", ".lora_down")) alphas[prefix] = value except Exception: pass module_keys = set(downs.keys()) | set(ups.keys()) module_keys_list = list(module_keys) total_modules = len(module_keys_list) for i, prefix in enumerate(module_keys_list, start=1): _progress(progress_cb, "load.modules", i, total_modules, prefix) if prefix not in downs or prefix not in ups: continue down = downs[prefix] up = ups[prefix] if down.ndim == 4 or up.ndim == 4: is_conv = True k_h = down.shape[2] if down.ndim == 4 else 1 k_w = down.shape[3] if down.ndim == 4 else 1 in_hw = (k_h, k_w) else: is_conv = False in_hw = (1, 1) rank = down.shape[0] alpha = float(alphas.get(prefix, rank)) pairs[prefix] = LoraPair( down=down, up=up, alpha=alpha, is_conv=is_conv, in_hw=in_hw, ) return pairs, metadata def _is_float_tensor(t: torch.Tensor) -> bool: return t.is_floating_point() def _prefer_bf16_when_mixed(dtypes: List[torch.dtype]) -> torch.dtype: """ If there is any mixed dtype situation and no explicit compute dtype was given, we favor bf16. """ uniq = list({dt for dt in dtypes}) if len(uniq) == 1: return uniq[0] return torch.bfloat16 def _resolve_working_dtype( tensors: List[torch.Tensor], explicit_compute_dtype: Optional[torch.dtype], ) -> torch.dtype: float_dtypes = [t.dtype for t in tensors if _is_float_tensor(t)] if not float_dtypes: return explicit_compute_dtype or torch.bfloat16 if explicit_compute_dtype is not None: return explicit_compute_dtype return _prefer_bf16_when_mixed(float_dtypes) def _cast_pair(lp: LoraPair, dtype: torch.dtype) -> LoraPair: if lp.down.dtype == dtype and lp.up.dtype == dtype: return lp return LoraPair( down=lp.down.to(dtype=dtype), up=lp.up.to(dtype=dtype), alpha=lp.alpha, is_conv=lp.is_conv, in_hw=lp.in_hw, ) def _safe_dot(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """ Dot product that guarantees dtype/device alignment and uses fp32 for stability. """ if a.device != b.device: b = b.to(a.device) a32 = a.to(torch.float32) b32 = b.to(torch.float32) return a32 @ b32 def compute_delta(lp: LoraPair, compute_dtype: torch.dtype) -> torch.Tensor: """ Compute full weight-space delta produced by this LoRA module. - Inputs are aligned to compute_dtype. - Matmul is done in fp32 for stability. - Output delta is fp32. """ down = lp.down up = lp.up if down.dtype != compute_dtype: down = down.to(compute_dtype) if up.dtype != compute_dtype: up = up.to(compute_dtype) alpha = lp.alpha rank = down.shape[0] scale = alpha / max(rank, 1) if lp.is_conv: r, c_in, k_h, k_w = down.shape out_channels = up.shape[0] down_flat = down.view(r, c_in * k_h * k_w).to(torch.float32) up_flat = up.view(out_channels, r).to(torch.float32) delta_flat = (up_flat @ down_flat) * float(scale) delta = delta_flat.view(out_channels, c_in, k_h, k_w) return delta return (up.to(torch.float32) @ down.to(torch.float32)) * float(scale) def tensor_norm(t: torch.Tensor) -> float: return float(torch.norm(t).item()) def choose_auto_rank(singular_values: torch.Tensor, threshold: float) -> int: if singular_values.numel() == 0: return 0 s2 = singular_values * singular_values total = torch.sum(s2).item() if total <= 0.0: return int(singular_values.shape[0]) cumulative = torch.cumsum(s2, dim=0) target = threshold * total mask = cumulative >= target indices = torch.nonzero(mask, as_tuple=False) if indices.numel() == 0: return int(singular_values.shape[0]) r = int(indices[0, 0].item() + 1) return r def delta_to_svd_factors( delta: torch.Tensor, rank: int, is_conv: bool, auto_rank_threshold: float, out_dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: if is_conv: out_channels, c_in, k_h, k_w = delta.shape mat = delta.view(out_channels, c_in * k_h * k_w) else: out_channels, c_in = delta.shape mat = delta mat32 = mat.to(torch.float32) u, s, vh = torch.linalg.svd(mat32, full_matrices=False) if rank == -1: r = choose_auto_rank(s, auto_rank_threshold) elif rank <= 0: r = s.shape[0] else: r = min(rank, s.shape[0]) if r <= 0: if is_conv: return ( torch.zeros((0, 0, 0, 0), dtype=out_dtype, device=delta.device), torch.zeros((0, 0), dtype=out_dtype, device=delta.device), ) return ( torch.zeros((0, 0), dtype=out_dtype, device=delta.device), torch.zeros((0, 0), dtype=out_dtype, device=delta.device), ) u_r = u[:, :r] s_r = s[:r] vh_r = vh[:r, :] sqrt_s = torch.sqrt(s_r.clamp_min(0.0)) up = u_r * sqrt_s.unsqueeze(0) down = sqrt_s.unsqueeze(1) * vh_r if is_conv: down = down.view(r, c_in, k_h, k_w) up = up.view(out_channels, r, 1, 1) down = down.to(out_dtype) up = up.to(out_dtype) return down.contiguous(), up.contiguous() def merge_mode_add( loras: List[Tuple[Dict[str, LoraPair], float]], include_patterns: List[str], exclude_patterns: List[str], explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (add)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.add", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue present: List[Tuple[LoraPair, float]] = [] for pairs, weight in loras: if prefix in pairs: present.append((pairs[prefix], weight)) if not present: continue is_conv = present[0][0].is_conv in_hw = present[0][0].in_hw for lp, _w in present[1:]: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' between LoRAs") compute_dtype = _resolve_working_dtype( [p[0].down for p in present] + [p[0].up for p in present], explicit_compute_dtype, ) downs_list: List[torch.Tensor] = [] ups_list: List[torch.Tensor] = [] total_rank = 0 for lp, w in present: lp = _cast_pair(lp, compute_dtype) rank_i = lp.down.shape[0] scale_i = w * (lp.alpha / max(rank_i, 1)) downs_list.append(lp.down) ups_list.append(lp.up * float(scale_i)) total_rank += rank_i new_down = torch.cat(downs_list, dim=0) new_up = torch.cat(ups_list, dim=1) new_alpha = float(total_rank) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_block_mix_weighted( loras: List[Tuple[Dict[str, LoraPair], float]], method: str, rank_value: int, preset: str, recipe: str, concept_mix: float, style_mix: float, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) != 2: raise ValueError("block-mix weighted mode requires exactly two LoRAs (A and B).") method_norm = (method or "svd").lower() if method_norm not in ("stack", "svd"): raise ValueError("block-mix method must be 'stack' or 'svd'.") a_pairs, a_weight = loras[0] b_pairs, b_weight = loras[1] all_prefixes = set(a_pairs.keys()) | set(b_pairs.keys()) prefixes_sorted = sorted(all_prefixes) preset_norm = (preset or "auto").lower().replace("_", "-") if preset_norm == "auto": preset_norm = _infer_block_mix_preset(prefixes_sorted) c_mix = float(concept_mix) s_mix = float(style_mix) c_mix = 0.0 if c_mix < 0.0 else (1.0 if c_mix > 1.0 else c_mix) s_mix = 0.0 if s_mix < 0.0 else (1.0 if s_mix > 1.0 else s_mix) merged: Dict[str, LoraPair] = {} total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc=f"Merging (block-mix-weighted:{method_norm})", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, f"merge.block-mix-weighted.{method_norm}", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue role = _block_mix_role(prefix, preset_norm) mix_a = c_mix if role == "concept" else s_mix mix_b = 1.0 - mix_a present: List[Tuple[LoraPair, float]] = [] if prefix in a_pairs and mix_a != 0.0: present.append((a_pairs[prefix], float(a_weight) * float(mix_a))) if prefix in b_pairs and mix_b != 0.0: present.append((b_pairs[prefix], float(b_weight) * float(mix_b))) if not present: continue ref_lp = present[0][0] is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw for lp, _w in present[1:]: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' in block-mix weighted mode") if method_norm == "stack": compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) downs_list: List[torch.Tensor] = [] ups_list: List[torch.Tensor] = [] total_rank = 0 for lp, w in present: lp = _cast_pair(lp, compute_dtype) rank_i = lp.down.shape[0] scale_i = float(w) * (lp.alpha / max(rank_i, 1)) downs_list.append(lp.down) ups_list.append(lp.up * float(scale_i)) total_rank += int(rank_i) new_down = torch.cat(downs_list, dim=0) new_up = torch.cat(ups_list, dim=1) new_alpha = float(total_rank) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) continue compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) merged_delta = None for lp, w in present: d = compute_delta(lp, compute_dtype) * float(w) merged_delta = d if merged_delta is None else (merged_delta + d) if merged_delta is None: continue svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def block_mix_routing_report( a_pairs: Dict[str, LoraPair], b_pairs: Dict[str, LoraPair], preset: str, recipe: str, include_patterns: List[str], exclude_patterns: List[str], ) -> Dict[str, Any]: all_prefixes = set(a_pairs.keys()) | set(b_pairs.keys()) prefixes_sorted = sorted(all_prefixes) preset_norm = (preset or "auto").lower().replace("_", "-") if preset_norm == "auto": preset_norm = _infer_block_mix_preset(prefixes_sorted) recipe_norm = (recipe or "").lower().replace("-", "_") routed_a = 0 routed_b = 0 fallback_to_a = 0 fallback_to_b = 0 missing_both = 0 skipped_filter = 0 requested_a = 0 requested_b = 0 concept = 0 style = 0 for prefix in prefixes_sorted: if not module_is_included(prefix, include_patterns, exclude_patterns): skipped_filter += 1 continue role = _block_mix_role(prefix, preset_norm) if role == "style": style += 1 else: concept += 1 idx = _block_mix_pick_lora_index(prefix, preset_norm, recipe_norm) if idx == 0: requested_a += 1 if prefix in a_pairs: routed_a += 1 elif prefix in b_pairs: routed_b += 1 fallback_to_b += 1 else: missing_both += 1 else: requested_b += 1 if prefix in b_pairs: routed_b += 1 elif prefix in a_pairs: routed_a += 1 fallback_to_a += 1 else: missing_both += 1 included_total = (len(prefixes_sorted) - skipped_filter) return { "preset": preset_norm, "recipe": recipe_norm, "modules_total": len(prefixes_sorted), "modules_included": included_total, "modules_skipped_filter": skipped_filter, "requested_a": requested_a, "requested_b": requested_b, "routed_a": routed_a, "routed_b": routed_b, "fallback_to_a": fallback_to_a, "fallback_to_b": fallback_to_b, "missing_both": missing_both, "role_concept": concept, "role_style": style, } def _infer_block_mix_preset(prefixes: List[str]) -> str: for p in prefixes: if p.startswith("diffusion_model.layers."): return "zimg-turbo" if p.startswith("lora_unet_double_blocks_") or p.startswith("lora_unet_single_blocks_"): return "flux" if p.startswith("lora_unet_blocks_"): return "wan" if p.startswith("transformer_blocks."): return "qwen" if p.startswith("lora_te1_") or p.startswith("lora_te2_"): return "sdxl" if p.startswith("lora_unet_"): return "sd" return "generic" def _block_mix_role(prefix: str, preset: str) -> str: p = prefix.lower() if preset == "zimg-turbo": if ".attention." in p: return "concept" if ".feed_forward." in p or ".adaln_modulation." in p: return "style" return "concept" if preset == "flux": if "_attn_" in p: return "concept" if "_mlp_" in p or "_mod_" in p or "_mod_lin" in p: return "style" return "concept" if preset == "wan": if "_attn_" in p: return "concept" if "_ffn_" in p or "_mlp_" in p: return "style" return "concept" if preset == "qwen": if ".attn." in p: return "concept" if ".mlp." in p or "_mlp" in p: return "style" return "concept" attn_markers = ( "attn", "to_q", "to_k", "to_v", "to_out", "q_proj", "k_proj", "v_proj", "out_proj", "proj", ) if any(m in p for m in attn_markers): return "concept" style_markers = ("mlp", "ff", "feed_forward", "res", "conv", "adaln", "norm") if any(m in p for m in style_markers): return "style" return "concept" def _block_mix_pick_lora_index(prefix: str, preset: str, recipe: str) -> int: r = (recipe or "").lower().replace("-", "_") if r in ("all_a", "only_a"): return 0 if r in ("all_b", "only_b"): return 1 if preset == "flux" and r in ("img_a_txt_b", "img_b_txt_a"): p = prefix.lower() if "_img_" in p: return 0 if r == "img_a_txt_b" else 1 if "_txt_" in p: return 1 if r == "img_a_txt_b" else 0 return 0 role = _block_mix_role(prefix, preset) if r in ("concept_a_style_b", "concept_from_a_style_from_b", "attn_a_ffn_b"): return 0 if role == "concept" else 1 if r in ("concept_b_style_a", "concept_from_b_style_from_a", "attn_b_ffn_a"): return 1 if role == "concept" else 0 return 0 def merge_mode_block_mix( loras: List[Tuple[Dict[str, LoraPair], float]], method: str, rank_value: int, preset: str, recipe: str, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) != 2: raise ValueError("block-mix mode requires exactly two LoRAs (A and B).") method_norm = (method or "svd").lower() if method_norm not in ("stack", "svd"): raise ValueError("block-mix method must be 'stack' or 'svd'.") a_pairs, a_weight = loras[0] b_pairs, b_weight = loras[1] all_prefixes = set(a_pairs.keys()) | set(b_pairs.keys()) prefixes_sorted = sorted(all_prefixes) preset_norm = (preset or "auto").lower().replace("_", "-") if preset_norm == "auto": preset_norm = _infer_block_mix_preset(prefixes_sorted) merged: Dict[str, LoraPair] = {} total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc=f"Merging (block-mix:{method_norm})", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, f"merge.block-mix.{method_norm}", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue idx = _block_mix_pick_lora_index(prefix, preset_norm, recipe) present: List[Tuple[LoraPair, float]] = [] if idx == 0: if prefix in a_pairs: present.append((a_pairs[prefix], float(a_weight))) elif prefix in b_pairs: present.append((b_pairs[prefix], float(b_weight))) else: if prefix in b_pairs: present.append((b_pairs[prefix], float(b_weight))) elif prefix in a_pairs: present.append((a_pairs[prefix], float(a_weight))) if not present: continue ref_lp = present[0][0] is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw for lp, _w in present[1:]: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' in block-mix mode") if method_norm == "stack": compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) downs_list: List[torch.Tensor] = [] ups_list: List[torch.Tensor] = [] total_rank = 0 for lp, w in present: lp = _cast_pair(lp, compute_dtype) rank_i = lp.down.shape[0] scale_i = float(w) * (lp.alpha / max(rank_i, 1)) downs_list.append(lp.down) ups_list.append(lp.up * float(scale_i)) total_rank += int(rank_i) new_down = torch.cat(downs_list, dim=0) new_up = torch.cat(ups_list, dim=1) new_alpha = float(total_rank) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) continue compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) merged_delta = None for lp, w in present: d = compute_delta(lp, compute_dtype) * float(w) merged_delta = d if merged_delta is None else (merged_delta + d) if merged_delta is None: continue svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_add_diff( loras: List[Tuple[Dict[str, LoraPair], float]], rank_value: int, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) < 2: raise ValueError("add-diff mode requires at least two LoRAs (base + one or more others).") base_pairs, base_weight = loras[0] all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (add-diff)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.add-diff", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue base_lp = base_pairs.get(prefix, None) others: List[Tuple[LoraPair, float]] = [] for pairs, w in loras[1:]: if prefix in pairs: others.append((pairs[prefix], w)) if base_lp is None and not others: continue ref_lp = base_lp if base_lp is not None else others[0][0] is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw if base_lp is not None and (base_lp.is_conv != is_conv or base_lp.in_hw != in_hw): raise ValueError(f"Incompatible shapes for module '{prefix}' in add-diff (base).") for lp, _w in others: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' in add-diff (other).") compute_dtype = _resolve_working_dtype( ([base_lp.down, base_lp.up] if base_lp is not None else []) + [t for lp, _w in others for t in (lp.down, lp.up)], explicit_compute_dtype, ) if base_lp is not None: delta_base_raw = compute_delta(base_lp, compute_dtype) delta_base = delta_base_raw * float(base_weight) else: delta_base_raw = None delta_base = torch.zeros_like(compute_delta(ref_lp, compute_dtype)) merged_delta = delta_base.clone() for lp_other, weight_other in others: delta_other = compute_delta(lp_other, compute_dtype) if delta_base_raw is None: merged_delta = merged_delta + float(weight_other) * delta_other else: merged_delta = merged_delta + float(weight_other) * (delta_other - delta_base_raw) svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_svd( loras: List[Tuple[Dict[str, LoraPair], float]], rank_value: int, preserve_norm: bool, cap_mult: Optional[float], include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (svd)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.svd", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue present: List[Tuple[LoraPair, float]] = [] for pairs, weight in loras: if prefix in pairs: present.append((pairs[prefix], weight)) if not present: continue ref_lp = present[0][0] is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw for lp, _w in present[1:]: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' between LoRAs") compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) deltas: List[torch.Tensor] = [] norms: List[float] = [] for lp, w in present: d = compute_delta(lp, compute_dtype) deltas.append(d * float(w)) norms.append(tensor_norm(d)) base_norm = sum(norms) / max(len(norms), 1) merged_delta = torch.zeros_like(deltas[0]) for d in deltas: merged_delta = merged_delta + d if preserve_norm: merged_norm = tensor_norm(merged_delta) if merged_norm > 1e-8 and base_norm > 1e-8: merged_delta = merged_delta * (base_norm / merged_norm) if cap_mult is not None and cap_mult > 0.0 and base_norm > 0.0: merged_norm = tensor_norm(merged_delta) max_norm = base_norm * cap_mult if merged_norm > max_norm and merged_norm > 0.0: merged_delta = merged_delta * (max_norm / merged_norm) svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_add_orth( loras: List[Tuple[Dict[str, LoraPair], float]], rank_value: int, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) < 2: raise ValueError("add-orth mode requires at least two LoRAs (base + one or more others).") base_pairs, base_weight = loras[0] all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (add-orth)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.add-orth", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue base_lp = base_pairs.get(prefix, None) present_others: List[Tuple[LoraPair, float]] = [] for pairs, weight in loras[1:]: if prefix in pairs: present_others.append((pairs[prefix], weight)) if base_lp is None and not present_others: continue if base_lp is None: temp_loras = [(pairs, weight) for (pairs, weight) in loras if prefix in pairs] temp_pairs_list = [] for pairs, weight in temp_loras: temp_pairs_list.append(({prefix: pairs[prefix]}, weight)) return_fallback = merge_mode_add(temp_pairs_list, include_patterns, exclude_patterns, explicit_compute_dtype) merged.update(return_fallback) continue is_conv = base_lp.is_conv in_hw = base_lp.in_hw for lp, _w in present_others: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' in add-orth mode") compute_dtype = _resolve_working_dtype( [base_lp.down, base_lp.up] + [t for lp, _w in present_others for t in (lp.down, lp.up)], explicit_compute_dtype, ) delta_base = compute_delta(base_lp, compute_dtype) * float(base_weight) base_flat = delta_base.reshape(-1) base_norm_sq = float(_safe_dot(base_flat, base_flat).item()) merged_delta = delta_base.clone() eps = 1e-12 for lp_other, weight_other in present_others: delta_other = compute_delta(lp_other, compute_dtype) * float(weight_other) other_flat = delta_other.reshape(-1) if base_norm_sq < eps: merged_delta = merged_delta + delta_other continue coeff = float(_safe_dot(other_flat, base_flat).item()) / (base_norm_sq + eps) orth_flat = other_flat - coeff * base_flat orth_delta = orth_flat.view_as(delta_other) merged_delta = merged_delta + orth_delta svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_diff_export( loras: List[Tuple[Dict[str, LoraPair], float]], rank_value: int, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) < 2: raise ValueError("diff-export mode requires at least two LoRAs (base and capability).") base_pairs, base_weight = loras[0] cap_pairs, cap_weight = loras[1] all_prefixes = set(base_pairs.keys()) | set(cap_pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (diff-export)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.diff-export", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue base_lp = base_pairs.get(prefix, None) cap_lp = cap_pairs.get(prefix, None) if base_lp is None and cap_lp is None: continue ref_lp = cap_lp if cap_lp is not None else base_lp is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw if base_lp is not None and (base_lp.is_conv != is_conv or base_lp.in_hw != in_hw): raise ValueError(f"Incompatible shapes for module '{prefix}' in base LoRA (diff-export).") if cap_lp is not None and (cap_lp.is_conv != is_conv or cap_lp.in_hw != in_hw): raise ValueError(f"Incompatible shapes for module '{prefix}' in capability LoRA (diff-export).") compute_dtype = _resolve_working_dtype( ([base_lp.down, base_lp.up] if base_lp is not None else []) + ([cap_lp.down, cap_lp.up] if cap_lp is not None else []), explicit_compute_dtype, ) if base_lp is not None: delta_base = compute_delta(base_lp, compute_dtype) * float(base_weight) else: delta_base = torch.zeros_like(compute_delta(ref_lp, compute_dtype)) if cap_lp is not None: delta_cap = compute_delta(cap_lp, compute_dtype) * float(cap_weight) else: delta_cap = torch.zeros_like(compute_delta(ref_lp, compute_dtype)) delta_diff = delta_cap - delta_base if torch.allclose(delta_diff, torch.zeros_like(delta_diff)): continue svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( delta_diff, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_moe( loras: List[Tuple[Dict[str, LoraPair], float]], rank_value: int, moe_temperature: float, moe_hard: bool, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) < 1: raise ValueError("moe mode requires at least one LoRA.") base_pairs, base_weight = loras[0] all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (moe)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.moe", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue base_lp = base_pairs.get(prefix, None) experts: List[Tuple[LoraPair, float]] = [] for pairs, weight in loras[1:]: if prefix in pairs: experts.append((pairs[prefix], weight)) if base_lp is None and not experts: continue ref_lp = base_lp if base_lp is not None else (experts[0][0] if experts else None) if ref_lp is None: continue is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw if base_lp is not None and (base_lp.is_conv != is_conv or base_lp.in_hw != in_hw): raise ValueError(f"Incompatible shapes for module '{prefix}' in base LoRA (moe).") for lp, _w in experts: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' in moe experts.") compute_dtype = _resolve_working_dtype( ([base_lp.down, base_lp.up] if base_lp is not None else []) + [t for lp, _w in experts for t in (lp.down, lp.up)], explicit_compute_dtype, ) if base_lp is not None: delta_base = compute_delta(base_lp, compute_dtype) * float(base_weight) else: delta_base = torch.zeros_like(compute_delta(ref_lp, compute_dtype)) if not experts: merged_delta = delta_base else: scores: List[float] = [] deltas_expert: List[torch.Tensor] = [] for lp, weight in experts: d = compute_delta(lp, compute_dtype) * float(weight) deltas_expert.append(d) scores.append(tensor_norm(d)) scores_tensor = torch.tensor(scores, dtype=torch.float32, device=delta_base.device) if moe_hard: best_index = int(torch.argmax(scores_tensor).item()) gate = torch.zeros_like(scores_tensor) gate[best_index] = 1.0 else: temp = max(moe_temperature, 1e-6) scaled = scores_tensor / temp gate = torch.softmax(scaled, dim=0) merged_delta = delta_base.clone() for g, d in zip(gate, deltas_expert): merged_delta = merged_delta + float(g.item()) * d svd_rank = rank_value if svd_rank == 0: svd_rank = -1 new_down, new_up = delta_to_svd_factors( merged_delta, rank=svd_rank, is_conv=is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_obfuscate( loras: List[Tuple[Dict[str, LoraPair], float]], include_patterns: List[str], exclude_patterns: List[str], explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: if len(loras) < 1: raise ValueError("obfuscate mode requires at least one LoRA.") all_prefixes = set() for pairs, _w in loras: all_prefixes.update(pairs.keys()) merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(all_prefixes) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Merging (obfuscate)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.obfuscate", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue present: List[Tuple[LoraPair, float]] = [] for pairs, weight in loras: if prefix in pairs: present.append((pairs[prefix], weight)) if not present: continue ref_lp = present[0][0] is_conv = ref_lp.is_conv in_hw = ref_lp.in_hw device = ref_lp.down.device for lp, _w in present[1:]: if lp.is_conv != is_conv or lp.in_hw != in_hw: raise ValueError(f"Incompatible shapes for module '{prefix}' between LoRAs in obfuscate mode") compute_dtype = _resolve_working_dtype( [t for lp, _w in present for t in (lp.down, lp.up)], explicit_compute_dtype, ) down_blocks: List[torch.Tensor] = [] up_blocks_flat: List[torch.Tensor] = [] total_rank = 0 c_in = None k_h = None k_w = None for lp, w in present: lp = _cast_pair(lp, compute_dtype) r_i = lp.down.shape[0] if r_i == 0: continue scale_i = float(w) * (lp.alpha / max(r_i, 1)) if is_conv: down_i = lp.down.to(torch.float32) up_i = lp.up.to(torch.float32) c_in = down_i.shape[1] k_h = down_i.shape[2] k_w = down_i.shape[3] up_i_flat = up_i.view(up_i.shape[0], r_i) up_i_flat = up_i_flat * scale_i down_blocks.append(down_i) up_blocks_flat.append(up_i_flat) else: down_i = lp.down.to(torch.float32) up_i = (lp.up.to(torch.float32) * scale_i) down_blocks.append(down_i) up_blocks_flat.append(up_i) total_rank += r_i if total_rank == 0: continue if is_conv: down_cat = torch.cat(down_blocks, dim=0) up_cat_flat = torch.cat(up_blocks_flat, dim=1) r_total = down_cat.shape[0] rand = torch.randn((r_total, r_total), device=device, dtype=torch.float32) q, _ = torch.linalg.qr(rand) down_flat = down_cat.view(r_total, -1) down_flat = q @ down_flat down_new32 = down_flat.view(r_total, c_in, k_h, k_w) up_cat_flat = up_cat_flat @ q.T up_new32 = up_cat_flat.view(up_cat_flat.shape[0], r_total, 1, 1) else: down_cat = torch.cat(down_blocks, dim=0) up_cat = torch.cat(up_blocks_flat, dim=1) r_total = down_cat.shape[0] rand = torch.randn((r_total, r_total), device=device, dtype=torch.float32) q, _ = torch.linalg.qr(rand) down_new32 = q @ down_cat up_new32 = up_cat @ q.T down_new = down_new32.to(compute_dtype) up_new = up_new32.to(compute_dtype) alpha_new = float(r_total) merged[prefix] = LoraPair( down=down_new.contiguous(), up=up_new.contiguous(), alpha=alpha_new, is_conv=is_conv, in_hw=in_hw, ) return merged def merge_mode_rebase( lora: Tuple[Dict[str, LoraPair], float], rank_value: int, include_patterns: List[str], exclude_patterns: List[str], auto_rank_threshold: float, explicit_compute_dtype: Optional[torch.dtype], progress_cb: Optional[ProgressCallback] = None, ) -> Dict[str, LoraPair]: pairs, weight = lora merged: Dict[str, LoraPair] = {} prefixes_sorted = sorted(pairs.keys()) total_modules = len(prefixes_sorted) for i, prefix in enumerate( tqdm(prefixes_sorted, desc="Rebasing (rebase)", unit="module", disable=progress_cb is not None), start=1, ): _progress(progress_cb, "merge.rebase", i, total_modules, prefix) if not module_is_included(prefix, include_patterns, exclude_patterns): continue lp = pairs[prefix] compute_dtype = _resolve_working_dtype([lp.down, lp.up], explicit_compute_dtype) delta = compute_delta(lp, compute_dtype) * float(weight) original_rank = lp.down.shape[0] target_rank = -1 if rank_value == 0 else rank_value new_down, new_up = delta_to_svd_factors( delta, rank=target_rank, is_conv=lp.is_conv, auto_rank_threshold=auto_rank_threshold, out_dtype=compute_dtype, ) new_rank = new_down.shape[0] if new_rank > original_rank: if lp.is_conv: new_down = new_down[:original_rank].contiguous() new_up = new_up[:, :original_rank, ...].contiguous() else: new_down = new_down[:original_rank].contiguous() new_up = new_up[:, :original_rank].contiguous() new_alpha = float(new_down.shape[0]) merged[prefix] = LoraPair( down=new_down, up=new_up, alpha=new_alpha, is_conv=lp.is_conv, in_hw=lp.in_hw, ) return merged def build_state_dict( merged: Dict[str, LoraPair], metadata_ref: Dict[str, str], dtype: torch.dtype, device: torch.device, progress_cb: Optional[ProgressCallback] = None, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, str]]: tensors: Dict[str, torch.Tensor] = {} key_style = _get_save_key_style(metadata_ref) items = list(merged.items()) total_modules = len(items) for i, (prefix, lp) in enumerate(items, start=1): _progress(progress_cb, "build.state_dict", i, total_modules, prefix) down = lp.down.to(dtype=dtype, device=device) up = lp.up.to(dtype=dtype, device=device) alpha_value = torch.tensor(lp.alpha, dtype=torch.float32, device=device) tensors[_down_key(prefix, key_style)] = down tensors[_up_key(prefix, key_style)] = up tensors[get_alpha_key(prefix)] = alpha_value meta = dict(metadata_ref or {}) meta.setdefault("format", "kohya-lora") return tensors, meta def summarize_pairs( pairs: Dict[str, LoraPair], weight: float, explicit_compute_dtype: Optional[torch.dtype], include_patterns: List[str], exclude_patterns: List[str], max_modules: int = 25, ) -> Dict[str, float]: prefixes = [p for p in sorted(pairs.keys()) if module_is_included(p, include_patterns, exclude_patterns)] if not prefixes: return { "modules": 0.0, "avg_rank": 0.0, "avg_alpha": 0.0, "avg_scale": 0.0, "mean_delta_norm": 0.0, "max_delta_norm": 0.0, } norms: List[float] = [] ranks: List[float] = [] alphas: List[float] = [] scales: List[float] = [] for prefix in prefixes[: max_modules]: lp = pairs[prefix] compute_dtype = _resolve_working_dtype([lp.down, lp.up], explicit_compute_dtype) delta = compute_delta(lp, compute_dtype) norms.append(tensor_norm(delta) * float(abs(weight))) r = float(lp.down.shape[0]) a = float(lp.alpha) ranks.append(r) alphas.append(a) scales.append(float(abs(weight)) * (a / max(r, 1.0))) return { "modules": float(len(prefixes)), "avg_rank": float(sum(ranks) / max(len(ranks), 1)), "avg_alpha": float(sum(alphas) / max(len(alphas), 1)), "avg_scale": float(sum(scales) / max(len(scales), 1)), "mean_delta_norm": float(sum(norms) / max(len(norms), 1)), "max_delta_norm": float(max(norms) if norms else 0.0), } def main(): parser = argparse.ArgumentParser( description=( "Comprehensive LoRA merge tool with modes: " "svd, rebase, add, add-diff, add-orth, diff-export, moe, obfuscate." ) ) parser.add_argument( "inputs", nargs="+", help="Input LoRAs as 'path[@weight]'. Example: A.safetensors@1.0 B.safetensors@-0.5", ) parser.add_argument("--out", required=True, help="Output .safetensors path") parser.add_argument( "--mode", type=str, default="svd", choices=["svd", "rebase", "add", "add-diff", "add-orth", "diff-export", "moe", "obfuscate", "block-mix"], help=( "Merge mode: " "svd = SVD rank-compressed merge, " "rebase = single-LoRA SVD rank-compression, " "add = exact linear stack (matches Comfy stacking behavior), " "add-diff = base + weighted differences toward others, " "add-orth = base + orthogonalized contributions, " "diff-export = pure difference LoRA, " "moe = mixture-of-experts per module (base + gated experts), " "obfuscate = stack-equivalent factor-space rebasis without SVD, " "block-mix = preset router that routes modules to LoRA A or B and merges via stack or svd." ), ) parser.add_argument( "--rank", type=int, default=64, help=( "Rank for SVD-based modes. " "Use 0 to trigger auto rank in svd/add-orth/diff-export/moe/add-diff/rebase." ), ) parser.add_argument( "--auto-rank-threshold", type=float, default=0.99, help="Energy ratio threshold for auto rank in SVD-based modes.", ) parser.add_argument( "--preserve-norm", action="store_true", help="In svd mode, preserve average per-module norm.", ) parser.add_argument( "--cap-mult", type=float, default=None, help="In svd mode, cap merged norm at cap_mult * mean(source_norms).", ) parser.add_argument( "--dtype", type=str, default="fp16", help="Output dtype: fp16, fp32, bf16", ) parser.add_argument( "--compute-dtype", type=str, default="auto", help=( "Internal merge dtype alignment: auto, bf16, fp16, fp32. " "auto => if mixed dtypes encountered, favor bf16." ), ) parser.add_argument( "--cpu", action="store_true", help="Force CPU even if CUDA is available", ) parser.add_argument( "--include-pattern", action="append", default=None, help="Only merge modules whose prefix contains this substring. Can be used multiple times.", ) parser.add_argument( "--exclude-pattern", action="append", default=None, help="Exclude modules whose prefix contains this substring. Can be used multiple times.", ) parser.add_argument( "--moe-temperature", type=float, default=1.0, help="Softmax temperature for moe mode (lower = sharper gating).", ) parser.add_argument( "--moe-hard", action="store_true", help="Use hard gating in moe mode (pick single best expert per module).", ) parser.add_argument( "--block-mix-method", type=str, default="svd", choices=["svd", "stack"], help="block-mix only: svd (delta merge then SVD) or stack (exact rank stacking).", ) parser.add_argument( "--block-mix-preset", type=str, default="auto", choices=["auto", "zimg-turbo", "flux", "wan", "qwen", "sd", "sdxl", "generic"], help="block-mix only: routing preset family.", ) parser.add_argument( "--block-mix-recipe", type=str, default="concept_a_style_b", help=( "block-mix only: recipe string. Examples: concept_a_style_b, concept_b_style_a, " "attn_a_ffn_b, img_a_txt_b, img_b_txt_a." ), ) parser.add_argument( "--report", action="store_true", help="Print a small diagnostic summary (module counts, ranks, scaling, delta norms).", ) args = parser.parse_args() weighted_paths = parse_weighted_paths(args.inputs) if not weighted_paths: raise SystemExit("No input LoRA files provided.") device = get_device(args.cpu) dtype = get_dtype(args.dtype) explicit_compute_dtype = get_compute_dtype(args.compute_dtype) include_patterns = args.include_pattern or [] exclude_patterns = args.exclude_pattern or [] loaded: List[Tuple[Dict[str, LoraPair], float]] = [] ref_meta: Dict[str, str] = {} for index, (path, weight) in enumerate(weighted_paths): pairs, meta = load_lora_pairs(path, device=device) print(f"Loaded {len(pairs)} modules from {path} (weight {weight})") if index == 0: ref_meta = meta loaded.append((pairs, weight)) if args.report: stats = summarize_pairs( pairs, weight=weight, explicit_compute_dtype=explicit_compute_dtype, include_patterns=include_patterns, exclude_patterns=exclude_patterns, ) print( "Report(input): " + json.dumps( { "path": path, "weight": weight, **stats, }, indent=2, ) ) if args.mode == "add": merged_pairs = merge_mode_add(loaded, include_patterns, exclude_patterns, explicit_compute_dtype) elif args.mode == "block-mix": merged_pairs = merge_mode_block_mix( loaded, method=args.block_mix_method, rank_value=args.rank, preset=args.block_mix_preset, recipe=args.block_mix_recipe, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "add-diff": merged_pairs = merge_mode_add_diff( loaded, rank_value=args.rank, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "add-orth": merged_pairs = merge_mode_add_orth( loaded, rank_value=args.rank, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "diff-export": merged_pairs = merge_mode_diff_export( loaded, rank_value=args.rank, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "moe": merged_pairs = merge_mode_moe( loaded, rank_value=args.rank, moe_temperature=args.moe_temperature, moe_hard=args.moe_hard, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "obfuscate": merged_pairs = merge_mode_obfuscate( loaded, include_patterns=include_patterns, exclude_patterns=exclude_patterns, explicit_compute_dtype=explicit_compute_dtype, ) elif args.mode == "rebase": if len(loaded) != 1: raise ValueError("rebase mode requires exactly one input LoRA.") merged_pairs = merge_mode_rebase( loaded[0], rank_value=args.rank, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) else: merged_pairs = merge_mode_svd( loaded, rank_value=args.rank, preserve_norm=args.preserve_norm, cap_mult=args.cap_mult if args.cap_mult is not None else None, include_patterns=include_patterns, exclude_patterns=exclude_patterns, auto_rank_threshold=args.auto_rank_threshold, explicit_compute_dtype=explicit_compute_dtype, ) state, meta = build_state_dict( merged_pairs, metadata_ref=ref_meta, dtype=dtype, device=device, ) if args.report: merged_stats = summarize_pairs( merged_pairs, weight=1.0, explicit_compute_dtype=explicit_compute_dtype, include_patterns=include_patterns, exclude_patterns=exclude_patterns, ) print("Report(merged): " + json.dumps(merged_stats, indent=2)) note_str = "Created by WAS Lora Merger" if args.mode == "obfuscate" else f"Created by WAS Lora Merger - Merging Mode: {args.mode}" meta.update( { "format": meta.get("format", "kohya-lora"), "merged_from": str(weighted_paths), "mode": args.mode, "rank": str(args.rank), "dtype": str(dtype).replace("torch.", ""), "compute_dtype": ( str(explicit_compute_dtype).replace("torch.", "") if explicit_compute_dtype is not None else "auto(bf16-if-mixed)" ), "preserve_norm": str(args.preserve_norm), "cap_mult": str(args.cap_mult) if args.cap_mult is not None else "None", "include_patterns": str(include_patterns), "exclude_patterns": str(exclude_patterns), "moe_temperature": str(args.moe_temperature), "moe_hard": str(args.moe_hard), "auto_rank_threshold": str(args.auto_rank_threshold), "tool": "WAS LoRA Merger", "note": note_str, } ) if args.mode == "obfuscate": meta.pop("mode", None) meta.pop("merged_from", None) meta.pop("rank", None) meta.pop("compute_dtype", None) meta.pop("include_patterns", None) meta.pop("exclude_patterns", None) meta.pop("moe_temperature", None) meta.pop("moe_hard", None) meta.pop("auto_rank_threshold", None) meta.pop("preserve_norm", None) meta.pop("cap_mult", None) out_dir = os.path.dirname(os.path.abspath(args.out)) or "." os.makedirs(out_dir, exist_ok=True) save_file(state, args.out, metadata=meta) print(f"Saved merged LoRA to: {args.out}") print(json.dumps(meta, indent=2)) if __name__ == "__main__": main()