convert.py

"""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 <input_stem>_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()