"""MiniMax-H3 Turbo LoRA — 4-step text-to-audio-video generation.

A lightweight LoRA that lets MiniMax-H3 render joint video + stereo audio in
**4 sampling steps** instead of the usual ~20, at a fraction of the wall-clock
cost. This single file is a self-contained generator: it loads the base H3 DiT
plus this LoRA, encodes the prompt with the Qwen3-VL text encoder, runs the
model's native dual-schedule sampler for 4 steps, decodes both streams and muxes
a playable mp4.

The audio stream runs on its own shifted flow schedule (video shift 12, audio
shift 3); each stream is integrated on its own clock, which is the schedule
semantics MiniMax-H3 was designed around. That is the only non-obvious part of
sampling — everything else is a plain Euler flow sampler.

Dependencies (see requirements.txt), plus a ComfyUI checkout for the H3 model /
VAE / text-encoder module definitions:

  git clone https://github.com/comfyanonymous/ComfyUI
  cd ComfyUI && git checkout 14b05228cef127ce529bc0c08660770d4af3e9a8

Base weights come from the official MiniMax-H3 release
(Comfy-Org/MiniMax-H3 on the Hugging Face Hub): the bf16 DiT, the int8 Qwen3-VL
text encoder, and the video + audio VAEs.

Usage:
  python generate.py \
    --comfyui /path/to/ComfyUI \
    --base   models/diffusion_models/minimax_h3_fl2va_bf16.safetensors \
    --lora   minimax_h3_turbo_4step.safetensors \
    --te     models/text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors \
    --video-vae models/vae/minimax_h3_video_vae_fp16.safetensors \
    --audio-vae models/vae/minimax_h3_audio_vae_fp32.safetensors \
    --prompt "A corgi in a tiny chef hat flipping a pancake, sizzling sounds." \
    --width 1344 --height 768 --frames 124 --out corgi.mp4

`minimax_h3_turbo_4step.safetensors` is the trained LoRA; the accompanying
`minimax_h3_turbo_4step_ema.safetensors` is a time-averaged variant — try both,
the trained one tends to be crisper on fast motion, the averaged one smoother.
"""

import argparse
import math
import os
import subprocess
import sys
import time
import wave

import torch
import torch.nn.functional as F


def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# ======================================================================
# Flow schedule (video shift 12 / audio shift 3, closed-form dual mapping)
# ======================================================================
SHIFT_VIDEO = 12.0
SHIFT_AUDIO = 3.0


def shift_sigma(u, shift):
    return shift * u / (1.0 + (shift - 1.0) * u)


def time_shift_sigma(sigma, from_shift, to_shift):
    base = sigma / (from_shift + sigma * (1.0 - from_shift))
    return to_shift * base / (1.0 + (to_shift - 1.0) * base)


def time_shift_slope(sigma, from_shift, to_shift):
    base = sigma / (from_shift + sigma * (1.0 - from_shift))
    return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / (
        from_shift * (1.0 + (to_shift - 1.0) * base) ** 2)


def timesteps(n, shift=SHIFT_VIDEO):
    """n-step video sigma grid: ts[0]=1 (pure noise) > ... > ts[n]=0."""
    return [shift_sigma(1.0 - i / n, shift) for i in range(n + 1)]


def audio_sigma(sigma_v):
    return time_shift_sigma(sigma_v, SHIFT_VIDEO, SHIFT_AUDIO)


def audio_slope(sigma_v):
    return time_shift_slope(sigma_v, SHIFT_VIDEO, SHIFT_AUDIO)


@torch.no_grad()
def sample(vfn, xv, xa, ts):
    """4-step Euler on the joint flow. The model returns the audio velocity
    already scaled by d(sigma_a)/d(sigma_v), so video steps on its own sigma
    delta while audio steps on its own schedule's delta (recovering the raw
    audio velocity by dividing out the slope). This dual-clock stepping is the
    schedule MiniMax-H3 expects; a single flat step on the video clock would
    over/under-shoot the audio stream badly at 4 steps.
    """
    for i in range(len(ts) - 1):
        ov, oa = vfn(xv, xa, ts[i])
        hv = ts[i + 1] - ts[i]
        sl = audio_slope(max(ts[i], 1e-6))
        ha = audio_sigma(ts[i + 1]) - audio_sigma(ts[i])
        xv = xv + hv * ov
        xa = xa + ha * (oa / sl)
    return xv, xa


# ======================================================================
# Functional forward (out-of-place, mirrors the reference module math)
# ======================================================================
def _rms(x, weight, eps):
    return F.rms_norm(x, (x.shape[-1],), weight, eps)


def _attn(attn, x, rope_cos, rope_sin):
    s = x.shape[0]
    heads, hd = attn.heads, attn.head_dim
    q, k, v = attn.qkv_proj(x).split(heads * hd, dim=-1)
    q = _rms(q.view(s, heads, hd), attn.q_norm.weight, attn.q_norm.eps)
    k = _rms(k.view(s, heads, hd), attn.k_norm.weight, attn.k_norm.eps)
    v = v.view(s, heads, hd)
    if rope_cos is not None:
        c, si = rope_cos[:, None, :], rope_sin[:, None, :]

        def rot(t):
            t96 = t[..., :96].float()
            x1, x2 = t96[..., :48], t96[..., 48:]
            return torch.cat([(x1 * c - x2 * si).to(t.dtype),
                              (x1 * si + x2 * c).to(t.dtype),
                              t[..., 96:]], dim=-1)

        q, k = rot(q), rot(k)
    q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v))
    out = F.scaled_dot_product_attention(q, k, v)
    return attn.out_proj(out.squeeze(0).transpose(0, 1).reshape(s, heads * hd))


def _mlp(mlp, x):
    x1, x2 = mlp.fc1(x).chunk(2, dim=-1)
    return mlp.fc2(F.silu(x1) * x2)


def _refiner(refiner, x):
    for blk in refiner.blocks:
        x = x + _attn(blk.attn, _rms(x, blk.norm1.weight, blk.norm1.eps),
                      None, None)
        x = x + _mlp(blk.mlp, _rms(x, blk.norm2.weight, blk.norm2.eps))
    return _rms(x, refiner.final_norm.weight, refiner.final_norm.eps)


def _apply_mod(h, shift, scale, segments):
    parts = []
    for a, b, row in segments:
        parts.append(h[a:b] * (1.0 + scale[row].to(h.dtype)) + shift[row].to(h.dtype))
    return torch.cat(parts)


def _apply_gate(x, gate, other, segments):
    parts = []
    for a, b, row in segments:
        parts.append(x[a:b] + other[a:b] * gate[row].to(x.dtype))
    return torch.cat(parts)


def _block(blk, h, mods, segments, rope_cos, rope_sin):
    sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = mods.unbind(dim=1)
    hn = _apply_mod(_rms(h, blk.norm1.weight, blk.norm1.eps), sh_msa, sc_msa, segments)
    h = _apply_gate(h, g_msa, _attn(blk.attn, hn, rope_cos, rope_sin), segments)
    hn = _apply_mod(_rms(h, blk.norm2.weight, blk.norm2.eps), sh_mlp, sc_mlp, segments)
    return _apply_gate(h, g_mlp, _mlp(blk.mlp, hn), segments)


class LoRALinear(torch.nn.Module):
    """Applies the low-rank update at run time in activation space:
    y = base(x) + B(A(x)). Folding it into the (bf16) base weight instead would
    round most of the update away when it is small relative to the weight, so we
    keep it as a separate matmul — same as how the update is meant to act."""

    def __init__(self, base, a, b):
        super().__init__()
        self.base = base
        self.a, self.b = a, b  # [rank, in], [out, rank]; alpha == rank -> scale 1

    def forward(self, x):
        return self.base(x) + F.linear(F.linear(x, self.a), self.b)


# ======================================================================
# Model load + LoRA (applied at run time, not merged)
# ======================================================================
def load_model(comfyui, base_path, lora_path, device, offload_adaln):
    import comfy.ldm.minimax.model as h3ref
    import comfy.ops
    import comfy.utils
    from safetensors.torch import load_file

    log(f"loading base DiT: {base_path}")
    sd = comfy.utils.load_torch_file(base_path)
    model = h3ref.MiniMaxH3Model(dtype=torch.bfloat16, device="cpu",
                                 operations=comfy.ops.disable_weight_init)
    missing, unexpected = model.load_state_dict(sd, strict=True, assign=True)
    assert not missing and not unexpected, (missing[:3], unexpected[:3])
    model.requires_grad_(False)
    model.eval()

    for i, blk in enumerate(model.blocks):
        blk.to(device)
    for mod in (model.token_refiner, model.final_layer, model.condition_proj,
                model.video_patch_proj, model.audio_patch_proj,
                model.time_embedder, model.rope):
        mod.to(device)

    log(f"applying LoRA: {lora_path}")
    lora = load_file(lora_path)
    names = sorted({k.rsplit(".lora_", 1)[0] for k in lora})
    # adaLN projections are read weight-first (bypassing their module), so their
    # LoRA can't ride a wrapper — stash it and add the delta where adaLN is built.
    model._adaln_lora = {}      # block index -> (a, b)
    model._final_adaln_lora = None
    n_wrap = 0
    for name in names:
        a = lora[name + ".lora_A.weight"].to(device, torch.bfloat16)
        b = lora[name + ".lora_B.weight"].to(device, torch.bfloat16)
        if name.endswith("adaln_proj.linear"):
            if name.startswith("final_layer"):
                model._final_adaln_lora = (a, b)
            else:
                model._adaln_lora[int(name.split(".")[1])] = (a, b)
        else:
            parent = model.get_submodule(name.rsplit(".", 1)[0])
            setattr(parent, name.rsplit(".", 1)[1],
                    LoRALinear(model.get_submodule(name), a, b))
            n_wrap += 1
    log(f"LoRA: {n_wrap} wrapped + {len(model._adaln_lora)} adaLN "
        f"+ {1 if model._final_adaln_lora else 0} final")

    if offload_adaln:
        # The per-layer adaLN projection is huge (2688 -> 96768) but depends only
        # on the timestep, of which there are a handful per denoise. Keep it in
        # CPU fp32 to save ~13 GB of VRAM; the matmul is cheap at 4 steps.
        for blk in model.blocks:
            lin = blk.adaln_proj.linear
            lin.weight.data = lin.weight.data.float().cpu()
            lin.bias.data = lin.bias.data.float().cpu()
    return model, h3ref


VISUAL_COND_T = 0.999


def timestep_rows(model, sigma_v):
    sigma_v = float(max(sigma_v, 1e-6))
    t_v = 1.0 - sigma_v
    t_a = 1.0 - time_shift_sigma(sigma_v, model.sigma_shift_video,
                                 model.sigma_shift_audio)
    seg_t = {"text": t_v, "video": t_v, "audio": t_a}
    unique_t = sorted({t_v, t_a})
    return seg_t, unique_t, {t: i for i, t in enumerate(unique_t)}


def adaln_mods(model, unique_t, device, offload, cache):
    key = tuple(round(t, 9) for t in unique_t)
    if key in cache:
        return cache[key]
    ts = torch.tensor(unique_t, dtype=torch.float32, device=device)
    with torch.no_grad():
        temb = model.time_embedder(ts).float()                 # [M, 2688] GPU
        si = F.silu(temb)
        si_base = si.cpu() if offload else si.to(torch.bfloat16)
        outs = torch.stack([F.linear(si_base, b.adaln_proj.linear.weight,
                                     b.adaln_proj.linear.bias)
                            for b in model.blocks])             # [50, M, 96768]
        mods = outs.to(device, torch.bfloat16)
        if getattr(model, "_adaln_lora", None):
            # run-time low-rank delta, on GPU (base built in CPU fp32 under offload)
            si_g = si.to(torch.bfloat16)
            for idx, (a, b) in model._adaln_lora.items():
                mods[idx] = mods[idx] + F.linear(F.linear(si_g, a), b)
    M, H = len(unique_t), model.hidden_size
    mods = mods.view(len(model.blocks), M, 3, 6, H).reshape(
        len(model.blocks), M * 3, 6, H)
    temb_bf = temb.to(torch.bfloat16)
    cache[key] = (mods, temb_bf)
    return mods, temb_bf


class Prepared:
    """Static packed-sequence structure for one (text_len, shape) signature."""

    def __init__(self, model, h3ref, text_len, video_shape, audio_t, tags,
                 device):
        _, _, lt, lh, lw = video_shape
        self.video_shape = tuple(video_shape)
        self.lat_pad = ((lh + 1) // 2 * 2, (lw + 1) // 2 * 2)
        self.layout = h3ref.PackedLayout(text_len, lt, *self.lat_pad, audio_t)
        pos = self.layout.position_ids.to(torch.float32).to(device)
        inv = model.rope.inv_freq.to(device)
        ang = (pos.unsqueeze(-1) * inv.view(1, 1, -1)).flatten(1)
        self.rope_cos, self.rope_sin = torch.cos(ang), torch.sin(ang)

        segs = []
        for a, b, kind in self.layout.segments:
            if kind == "text" and tags is not None:
                tg = tags.view(-1).tolist()
                run = 0
                for i in range(1, b - a + 1):
                    if i == b - a or tg[i] != tg[run]:
                        segs.append((a + run, a + i, int(tg[run]), kind))
                        run = i
            else:
                tag = {"text": 1, "video": 0, "audio": 2}[kind]
                segs.append((a, b, tag, kind))
        self.seg_template = segs
        (self.video_seg,) = [(a, b) for a, b, k in self.layout.segments if k == "video"]
        (self.audio_seg,) = [(a, b) for a, b, k in self.layout.segments if k == "audio"]


@torch.no_grad()
def forward(model, h3ref, prep, video_x, audio_x, sigma_v, context, device,
            offload, cache):
    """One denoise evaluation in the sigma_v domain. Returns
    (video_velocity, audio_velocity * slope), matching what the sampler wants."""
    import comfy.ldm.common_dit
    video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, model.patch_size)
    orig_t, orig_h, orig_w = prep.video_shape[2:]

    sigma_v = float(max(sigma_v, 1e-6))
    seg_t, unique_t, t_row = timestep_rows(model, sigma_v)
    segments = [(a, b, t_row[seg_t[k]] * 3 + tag)
                for a, b, tag, k in prep.seg_template]

    base_mods, t_emb = adaln_mods(model, unique_t, device, offload, cache)
    silu_temb = F.silu(t_emb)

    video_rows = h3ref.patchify_video(video_x.to(torch.float32), model.patch_size)
    audio_rows = h3ref.pack_audio(audio_x.to(torch.float32))
    video_embed = model.video_patch_proj(video_rows).to(torch.bfloat16)
    audio_embed = model.audio_patch_proj(audio_rows).to(torch.bfloat16)

    with torch.autocast("cuda", dtype=torch.bfloat16):
        text_states = context[0]
        if text_states.shape[-1] != model.hidden_size:
            text_states = _refiner(model.token_refiner,
                                   model.condition_proj(text_states))
        pieces = []
        for a, b, kind in prep.layout.segments:
            if kind == "text":
                pieces.append(text_states)
            elif kind == "video":
                pieces.append(video_embed)
            else:
                pieces.append(audio_embed)
        h = torch.cat(pieces)
        for i, blk in enumerate(model.blocks):
            h = _block(blk, h, base_mods[i], segments,
                       prep.rope_cos, prep.rope_sin)

    fl = model.final_layer
    with torch.autocast("cuda", dtype=torch.bfloat16):
        si_t = F.silu(t_emb)
        f_mod = fl.adaln_proj.linear(si_t)
        if getattr(model, "_final_adaln_lora", None):
            a, b = model._final_adaln_lora
            f_mod = f_mod + F.linear(F.linear(si_t, a), b)
    f_shift, f_scale = f_mod.view(len(unique_t), 2, model.hidden_size).unbind(1)
    (va, vb), (aa, ab) = prep.video_seg, prep.audio_seg
    vrow, arow = t_row[seg_t["video"]], t_row[seg_t["audio"]]
    hn = _rms(h, fl.norm.weight, fl.norm.eps)
    hv = (hn[va:vb] * (1.0 + f_scale[vrow]) + f_shift[vrow]).to(torch.float32)
    ha = (hn[aa:ab] * (1.0 + f_scale[arow]) + f_shift[arow]).to(torch.float32)
    v_rows, a_rows = fl.video_out(hv), fl.audio_out(ha)

    lt = video_x.shape[2]
    video_out = h3ref.unpatchify_video(v_rows, lt, prep.lat_pad[0] // 2,
                                       prep.lat_pad[1] // 2, model.latents_dim,
                                       model.patch_size)[:, :, :orig_t, :orig_h, :orig_w]
    audio_out = h3ref.unpack_audio(a_rows)
    slope_a = time_shift_slope(sigma_v, model.sigma_shift_video,
                               model.sigma_shift_audio)
    return -video_out.to(video_x.dtype), (-slope_a) * audio_out.to(audio_x.dtype)


# ======================================================================
# Text encode / decode / mux
# ======================================================================
def encode_prompt(comfyui, te_path, prompt, device):
    import comfy.model_management
    import comfy.sd
    log(f"loading text encoder: {te_path}")
    clip = comfy.sd.load_clip([te_path], clip_type=comfy.sd.CLIPType.MINIMAX)
    cond = clip.encode_from_tokens_scheduled(clip.tokenize(prompt))
    ca, ex = cond[0][0], cond[0][1]
    tags = ex.get("minimax_token_tags")
    ctx = ca.to(device, torch.bfloat16)
    tags = tags.to(device) if torch.is_tensor(tags) else tags
    del clip
    comfy.model_management.unload_all_models()
    comfy.model_management.soft_empty_cache()
    return ctx, tags


def _write_wav(path, waveform, sr):
    w = waveform.detach().cpu().float()
    if w.ndim == 3:
        w = w[0]
    w = w.clamp(-1.0, 1.0)
    ch = w.shape[0]
    pcm = (w.transpose(0, 1).contiguous().numpy() * 32767.0).astype("<i2")
    with wave.open(path, "wb") as f:
        f.setnchannels(ch)
        f.setsampwidth(2)
        f.setframerate(int(sr))
        f.writeframes(pcm.tobytes())


def save_mp4(images, waveform, sr, fps, out_path):
    import imageio.v2 as imageio
    import imageio_ffmpeg
    frames = images.detach().cpu().float().clamp(0, 1).mul(255).round().to(
        torch.uint8).numpy()
    tv, ta = out_path + ".v.mp4", out_path + ".a.wav"
    writer = imageio.get_writer(tv, fps=fps, codec="libx264", quality=8,
                                pixelformat="yuv420p", macro_block_size=1,
                                ffmpeg_log_level="error")
    for fr in frames:
        writer.append_data(fr)
    writer.close()
    _write_wav(ta, waveform, sr)
    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
    subprocess.run([ffmpeg, "-y", "-loglevel", "error", "-i", tv, "-i", ta,
                    "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest",
                    out_path], check=True)
    os.remove(tv)
    os.remove(ta)


# ======================================================================
# Main
# ======================================================================
def main():
    ap = argparse.ArgumentParser(description="MiniMax-H3 Turbo LoRA 4-step generator")
    ap.add_argument("--comfyui", required=True, help="path to a ComfyUI checkout @14b05228")
    ap.add_argument("--base", required=True, help="H3 bf16 DiT safetensors")
    ap.add_argument("--lora", required=True, help="turbo LoRA safetensors")
    ap.add_argument("--te", required=True, help="Qwen3-VL text encoder safetensors")
    ap.add_argument("--video-vae", required=True)
    ap.add_argument("--audio-vae", required=True)
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--out", default="out.mp4")
    ap.add_argument("--width", type=int, default=1344, help="multiple of 16 (canvas is 32-based)")
    ap.add_argument("--height", type=int, default=768)
    ap.add_argument("--frames", type=int, default=124, help="24 fps; snaps to the 17k+5 grid")
    ap.add_argument("--steps", type=int, default=4)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--offload-adaln", action="store_true",
                    help="keep the timestep-projection weights in CPU fp32 (saves ~13GB VRAM)")
    args = ap.parse_args()
    sys.path.insert(0, args.comfyui)  # ComfyUI supplies the H3 module definitions

    dev = "cuda"
    frames = args.frames
    while frames % 17 != 5:
        frames += 1
    lt = (frames - 5) // 17 * 5 + 2
    lh, lw = args.height // 16, args.width // 16
    audio_t = round(frames / 24 * 40)
    v_shape, a_shape = (1, 24, lt, lh, lw), (1, 32, 2, audio_t)
    ts = timesteps(args.steps)
    log(f"{args.width}x{args.height}x{frames}f ({frames/24:.1f}s) -> "
        f"video{v_shape} audio{a_shape}; {args.steps}-step grid "
        f"{['%.3f' % t for t in ts]}")

    ctx, tags = encode_prompt(args.comfyui, args.te, args.prompt, dev)
    model, h3ref = load_model(args.comfyui, args.base, args.lora, dev,
                              args.offload_adaln)
    prep = Prepared(model, h3ref, ctx.shape[1], v_shape, audio_t, tags, dev)

    g = torch.Generator(dev).manual_seed(args.seed)
    ga = torch.Generator(dev).manual_seed(args.seed + 1)
    nv = torch.randn(v_shape, generator=g, device=dev, dtype=torch.bfloat16)
    na = torch.randn(a_shape, generator=ga, device=dev, dtype=torch.bfloat16)

    cache = {}

    def vfn(xv, xa, sv):
        return forward(model, h3ref, prep, xv, xa, sv, ctx, dev,
                       args.offload_adaln, cache)

    log("sampling ...")
    t0 = time.time()
    with torch.inference_mode():
        zv, za = sample(vfn, nv, na, ts)
    log(f"sampled in {time.time()-t0:.1f}s")

    import comfy.sd
    import comfy.utils
    video_vae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(args.video_vae))
    audio_vae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(args.audio_vae))
    with torch.inference_mode():
        images = video_vae.decode(zv.float())
        if images.ndim == 5:
            images = images.reshape(-1, *images.shape[-3:])
        waveform = audio_vae.decode(za.float()).movedim(-1, 1)
        std = torch.std(waveform, dim=[1, 2], keepdim=True) * 5.0
        std[std < 1.0] = 1.0
        waveform = waveform / std
    sr = getattr(audio_vae, "audio_sample_rate_output",
                 getattr(audio_vae, "audio_sample_rate", 44100))
    save_mp4(images, waveform, sr, 24, args.out)
    log(f"done -> {args.out} ({os.path.getsize(args.out)/2**20:.1f}MB)")


if __name__ == "__main__":
    main()