h3_lora.py

"""Wushu Action LoRA support for the diffusers MiniMax-H3 transformer.

The Jojocodex wushu-action LoRA (`Jojocodex/minimax-h3-wushu-action-lora`) targets the ComfyUI reference module
tree — `diffusion_model.blocks.N.attn.qkv_proj`, `diffusion_model.blocks.N.mlp.fc1`, etc. — with rank 16 and no
alpha metadata (so scale = 1.0, matching the convention where alpha == rank). The shipped `_pruned` file carries
416 keys over 208 base modules (52 transformer blocks x {attn.qkv_proj, attn.out_proj, mlp.fc1, mlp.fc2}, the two
`token_refiner` blocks included) and has its `adaln_proj` rows removed (`__metadata__: {"adaln_pruned": "true"}`),
which is what makes it stackable with the Turbo acceleration LoRA.

The LoRA is applied by folding `scale * (lora_B @ lora_A)` into the bf16 weights rather than as runtime wrappers:
folding costs nothing per request and keeps the transformer a plain `nn.Module` for the ZeroGPU startup packing.
Deltas are computed in float32 and round once on the way back into bf16.

The diffusers conversion transforms:
- strip the `diffusion_model.` prefix the ai-toolkit export uses
- fused-QKV row thirds onto `attn.to_q/k/v`
- `SwiGLU` gate/value swap onto `ff.net.0.proj`
- `fc2` -> `ff.net.2`, `attn.out_proj` -> `attn.to_out.0`
- `blocks.` -> `transformer_blocks.`
- `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`
- `final_layer.adaln_proj.linear` -> `norm_out.linear` (not present in the pruned file)
"""

from __future__ import annotations

import os

import torch

# --- Wushu Action LoRA config ---
# The model card tells ComfyUI users to take the `_pruned` file; it is also the Turbo-compatible one.
WUSHU_LORA_REPO = os.environ.get("WUSHU_LORA_REPO", "Jojocodex/minimax-h3-wushu-action-lora")
WUSHU_LORA_FILE = os.environ.get("WUSHU_LORA_FILE", "wushu_action_h3_lora_v4_2000_pruned.safetensors")
# The card recommends strength 0.8~1.0 for the ComfyUI LoraLoader; folded at the top of that range.
WUSHU_LORA_STRENGTH = float(os.environ.get("WUSHU_LORA_STRENGTH", "1.0"))

# --- Turbo LoRA config (optional acceleration) ---
TURBO_LORA_REPO = os.environ.get("TURBO_LORA_REPO", "Comfy-Org/MiniMax-H3")
TURBO_LORA_FILE = os.environ.get(
    "TURBO_LORA_FILE", "loras/minimax_h3_fl2v_turbo_4step_v1.0_768p_comfyui_bf16.safetensors"
)
TURBO_LORA_STRENGTH = float(os.environ.get("TURBO_LORA_STRENGTH", "1.0"))

USE_TURBO = os.environ.get("USE_TURBO", "1").lower() not in ("0", "off", "none", "false")


def _targets(name: str, b: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
    """Map one reference-tree base name and its `lora_B` onto diffusers parameter key + row-transformed B."""
    name = name.removeprefix("diffusion_model.")

    if name.startswith("token_refiner.blocks."):
        target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
    elif name.startswith("blocks."):
        target = name.replace("blocks.", "transformer_blocks.", 1)
    else:
        target = name
    target = target.replace("final_layer.adaln_proj.linear", "norm_out.linear")

    if target.endswith(".attn.qkv_proj"):
        prefix = target.removesuffix("qkv_proj")
        return [
            (f"{prefix}to_{kind}.weight", part.contiguous())
            for kind, part in zip(("q", "k", "v"), b.split(inner_dim, dim=0))
        ]
    if target.endswith(".mlp.fc1"):
        gate, value = b.chunk(2, dim=0)
        return [(target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous())]
    if target.endswith(".mlp.fc2"):
        return [(target.replace(".mlp.fc2", ".ff.net.2") + ".weight", b)]
    if target.endswith(".attn.out_proj"):
        return [(target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", b)]
    # `adaln_proj.linear` (block-level and the final `norm_out.linear`): identical row layout on both sides.
    return [(target + ".weight", b)]


def _load_wushu_lora(inner_dim: int) -> dict:
    """Load the Jojocodex wushu action LoRA from the Hub."""
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    lora = load_file(hf_hub_download(WUSHU_LORA_REPO, WUSHU_LORA_FILE))
    bases = sorted({key.rsplit(".lora_", 1)[0] for key in lora})
    entries = []
    for name in bases:
        a = lora[f"{name}.lora_A.weight"]
        b = lora[f"{name}.lora_B.weight"]
        entries.extend((key, a, b_part) for key, b_part in _targets(name, b, inner_dim))
    return {
        "label": f"{WUSHU_LORA_REPO}/{WUSHU_LORA_FILE}",
        "scale": WUSHU_LORA_STRENGTH,  # alpha == rank, so the base scale is 1
        "entries": entries,
    }


def _load_turbo_lora(inner_dim: int) -> dict:
    """Load the MiniMax-H3 Turbo LoRA from Comfy-Org for 4-step accelerated inference.

    The Comfy-Org Turbo LoRA uses the kohya format with explicit `.alpha` keys per LoRA layer, so each entry's
    scale is `alpha / rank`, applied to `B` before the fused-QKV split.
    """
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    lora = load_file(hf_hub_download(TURBO_LORA_REPO, TURBO_LORA_FILE))

    all_keys = list(lora.keys())
    bases = sorted({key.rsplit(".lora_", 1)[0] if ".lora_" in key else key.rsplit(".alpha", 1)[0] for key in all_keys})

    entries = []
    for name in bases:
        a = lora[f"{name}.lora_A.weight"]
        b = lora[f"{name}.lora_B.weight"]
        alpha_key = f"{name}.alpha"
        if alpha_key in lora:
            scale = float(lora[alpha_key]) / a.shape[0]
        else:
            scale = 1.0
        b_scaled = b * (scale * TURBO_LORA_STRENGTH)
        entries.extend((key, a, b_part) for key, b_part in _targets(name, b_scaled, inner_dim))
    return {
        "label": f"{TURBO_LORA_REPO}/{TURBO_LORA_FILE}",
        "scale": 1.0,  # scale already applied per-entry above
        "entries": entries,
    }


def _apply(entries, params, sign: float) -> None:
    """Fold (sign * scale * (B @ A)) into each target parameter in place."""
    for key, a, b in entries:
        param = params.get(key)
        if param is None:
            raise KeyError(f"LoRA target `{key}` not found in the transformer")
        delta = sign * (b.to(torch.float32) @ a.to(torch.float32))
        param.data = (param.data.float() + delta.to(param.device)).to(param.dtype)


def apply_lora(transformer) -> str | None:
    """Fold the wushu action LoRA — and, when enabled, the Turbo LoRA — into the transformer weights.

    Returns a status line, or `None` when nothing could be folded.
    """
    inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
    params = dict(transformer.named_parameters())

    loaded = []

    try:
        wushu = _load_wushu_lora(inner_dim)
        _apply(wushu["entries"], params, wushu["scale"])
        loaded.append(f"wushu-action ({wushu['label']}, {len(wushu['entries'])} weights, scale={wushu['scale']})")
        print(f"[lora] wushu action LoRA folded: {len(wushu['entries'])} weights", flush=True)
    except Exception as error:
        print(f"[lora] WARNING: failed to load the wushu action LoRA: {error}", flush=True)

    if USE_TURBO:
        try:
            turbo = _load_turbo_lora(inner_dim)
            _apply(turbo["entries"], params, turbo["scale"])
            loaded.append(f"turbo ({turbo['label']}, {len(turbo['entries'])} weights)")
            print(f"[lora] turbo LoRA folded: {len(turbo['entries'])} weights", flush=True)
        except Exception as error:
            print(f"[lora] WARNING: failed to load the turbo LoRA: {error}", flush=True)

    if not loaded:
        return None

    return "LoRAs folded: " + " + ".join(loaded)