"""Convert a MiniMax-H3 Turbo LoRA (ComfyUI / generate.py layout) to Diffusers PEFT layout. The Turbo LoRA ships with ComfyUI module names and fused projections: * ``blocks.*.attn.qkv_proj`` — fused ``[q_all; k_all; v_all]`` (the in-memory layout after ComfyUI's load-time QKV reorder), not the raw checkpoint's per-head interleave. * ``blocks.*.mlp.fc1`` — fused ``[gate; value]``; Diffusers' ``SwiGLU`` wants ``[value; gate]``. * ``alpha == rank`` — no extra scale (``W_eff = W + B @ A``). This script renames everything onto ``MiniMaxH3Transformer3DModel``, splits the fused QKV LoRA into ``to_q`` / ``to_k`` / ``to_v`` (shared ``A``, split ``B``), and swaps the ``fc1`` halves so the low-rank update matches the converted base weights. Usage: ```bash python convert.py \ --input minimax_h3_turbo_4step_ckpt500.safetensors \ --output minimax_h3_turbo_4step_ckpt500_diffusers.safetensors ``` If ``--output`` is omitted, ``_diffusers`` is inserted before the ``.safetensors`` suffix. """ from __future__ import annotations import argparse from collections import defaultdict from pathlib import Path import torch from safetensors.torch import load_file, save_file # Must match MiniMaxH3Transformer3DModel / convert_minimax_h3_to_diffusers.py. NUM_ATTENTION_HEADS = 56 ATTENTION_HEAD_DIM = 128 INNER_DIM = NUM_ATTENTION_HEADS * ATTENTION_HEAD_DIM # 7168 FFN_DIM = 14336 PREFIX = "transformer" def _module_names(state_dict: dict[str, torch.Tensor]) -> list[str]: return sorted({key.rsplit(".lora_", 1)[0] for key in state_dict}) def _rename_base(name: str) -> str: """Map a ComfyUI module path (without ``.lora_*``) onto the Diffusers module path.""" if name.startswith("token_refiner.blocks."): name = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1) elif name.startswith("blocks."): name = name.replace("blocks.", "transformer_blocks.", 1) name = name.replace("final_layer.adaln_proj.linear", "norm_out.linear") name = name.replace(".attn.out_proj", ".attn.to_out.0") name = name.replace(".mlp.fc2", ".ff.net.2") name = name.replace(".mlp.fc1", ".ff.net.0.proj") return name def convert_lora_state_dict( src: dict[str, torch.Tensor], ) -> tuple[dict[str, torch.Tensor], dict[str, int]]: """Convert one Turbo LoRA state dict into Diffusers PEFT keys (with ``transformer.`` prefix).""" out: dict[str, torch.Tensor] = {} counts: dict[str, int] = defaultdict(int) for name in _module_names(src): a = src[f"{name}.lora_A.weight"] b = src[f"{name}.lora_B.weight"] if a.ndim != 2 or b.ndim != 2: raise ValueError(f"{name}: expected 2-D LoRA matrices, got A{tuple(a.shape)} B{tuple(b.shape)}") if a.shape[0] != b.shape[1]: raise ValueError(f"{name}: rank mismatch A{tuple(a.shape)} vs B{tuple(b.shape)}") # qkv: split fused [q;k;v] B into three LoRAs that share A. if name.endswith(".attn.qkv_proj"): if b.shape[0] != 3 * INNER_DIM: raise ValueError( f"{name}: fused qkv B has {b.shape[0]} rows, expected {3 * INNER_DIM} " f"(= 3 * {INNER_DIM})." ) base = _rename_base(name[: -len(".attn.qkv_proj")]) bq, bk, bv = b.split(INNER_DIM, dim=0) for suffix, b_part in (("to_q", bq), ("to_k", bk), ("to_v", bv)): key = f"{PREFIX}.{base}.attn.{suffix}" # Clone A so to_q/to_k/to_v do not share storage (safetensors forbids that). out[f"{key}.lora_A.weight"] = a.detach().clone().contiguous() out[f"{key}.lora_B.weight"] = b_part.detach().clone().contiguous() counts["qkv_split"] += 1 continue base = _rename_base(name) # fc1 / SwiGLU: reference stores [gate; value], Diffusers wants [value; gate]. if name.endswith(".mlp.fc1"): if b.shape[0] != 2 * FFN_DIM: raise ValueError( f"{name}: fc1 B has {b.shape[0]} rows, expected {2 * FFN_DIM} (= 2 * {FFN_DIM})." ) gate, value = b.chunk(2, dim=0) b = torch.cat([value, gate], dim=0).contiguous() counts["fc1_swap"] += 1 else: counts["rename"] += 1 key = f"{PREFIX}.{base}" out[f"{key}.lora_A.weight"] = a.detach().clone().contiguous() out[f"{key}.lora_B.weight"] = b.detach().clone().contiguous() return out, dict(counts) def network_alphas_from_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, float]: """``alpha == rank`` for every module (Turbo LoRA convention).""" alphas: dict[str, float] = {} for key, tensor in state_dict.items(): if key.endswith(".lora_B.weight") and tensor.ndim > 1: base = key[: -len(".lora_B.weight")] alphas[f"{base}.alpha"] = float(tensor.shape[1]) return alphas def default_output_path(input_path: str | Path) -> Path: path = Path(input_path) stem = path.stem if stem.endswith("_diffusers"): return path return path.with_name(f"{stem}_diffusers{path.suffix}") def main() -> None: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument( "--input", required=True, help="ComfyUI / generate.py Turbo LoRA safetensors " "(e.g. minimax_h3_turbo_4step_ckpt500.safetensors from larryvrh/MiniMax-H3-Turbo-Lora)", ) parser.add_argument( "--output", default=None, help="Diffusers PEFT LoRA safetensors (keys prefixed with transformer.). " "Defaults to _diffusers.safetensors", ) args = parser.parse_args() output = Path(args.output) if args.output is not None else default_output_path(args.input) print(f"loading {args.input}") src = load_file(args.input, device="cpu") dst, counts = convert_lora_state_dict(src) ranks = sorted({int(v.shape[1]) for k, v in dst.items() if k.endswith(".lora_B.weight")}) print( f"converted {len(src)} -> {len(dst)} tensors; " f"qkv_split={counts.get('qkv_split', 0)} " f"fc1_swap={counts.get('fc1_swap', 0)} " f"rename={counts.get('rename', 0)}; ranks={ranks}" ) # Keep the original dtype (bf16). Metadata is informational for humans / loaders. metadata = { "format": "pt", "base_model": "MiniMax-H3", "application": "W_eff = W + lora_B @ lora_A (alpha == rank)", "sampler_steps": "4", "converted_from": "larryvrh/MiniMax-H3-Turbo-Lora (ComfyUI layout)", } save_file(dst, str(output), metadata=metadata) print(f"saved {output}") # Print a tiny alpha hint so callers can wire network_alphas correctly. alphas = network_alphas_from_state_dict(dst) alpha_values = sorted(set(alphas.values())) print(f"network alphas (== rank): {alpha_values} across {len(alphas)} modules") if __name__ == "__main__": main()