app.py

"""MiniMax-H3 Wushu Action LoRA — text-to-video demo for the martial-arts / kung-fu motion LoRA.

This Space is the denoising half of a split MiniMax-H3 deployment:
- The 62 GiB Qwen3-VL text encoder runs in the conditioner Space (`multimodalart/qwen3vl-conditioner`),
  called over the gradio API for each request.
- This Space loads the 61.7 GiB transformer + 10.4 GiB VAEs (77.3 GB total) and runs the denoising loop
  and the video + audio decode on the GPU.
- The Jojocodex wushu action LoRA (rank 16, `_pruned`) is folded into the transformer weights at startup,
  adding human martial-arts motion: punches, kicks, combination forms, staff technique.
- The Comfy-Org MiniMax-H3 Turbo LoRA (4-step) is folded on top, which the LoRA card says is supported
  because the `_pruned` file has its `adaln_proj` rows removed.

The LoRA was trained with ai-toolkit on 455 curated wushu clips at 90 frames / 24 fps — so the demo's default
duration is 3.75 s, which is exactly the 90-frame window it saw.
"""

from __future__ import annotations

import math
import os
import tempfile
import time
import traceback
from functools import cache

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 72 GiB load can happen
# at startup rather than on GPU time.
import spaces
import gradio as gr

# --- Configuration ---
MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
LORA_REPO = "Jojocodex/minimax-h3-wushu-action-lora"

# `pack` places the transformer at startup, `lazy` moves everything on the first GPU call.
PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()
# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
# The transformer alone is 61.7 GiB, so the 48 GB `large` booking cannot hold it.
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# --- Canvas definitions (labels are the wire contract with the conditioner) ---
CANVASES = {
    # 16:9
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    # 9:16
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    # 1:1
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    # 4:3 / 3:4
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    # 21:9
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"

FPS = 24
FRAMES_PER_CHUNK = 17
LATENTS_PER_CHUNK = 5

MIN_UI_DURATION = 2.0
MAX_UI_DURATION = 14.0
# 3.75 s == 90 frames == 17 * 5 + 5, the clip length the LoRA was trained on.
TRAINED_DURATION = 3.75

DEFAULT_STEPS = 4
DEFAULT_SEED = 42

# The LoRA has no literal trigger token: the card says to activate it with a natural-language *action* description,
# and lists the four technique families it was captioned around. Each cue below is appended to the user's prompt so
# a request lands inside the family the user picked, in the wording the training captions used.
TECHNIQUES: dict[str, str] = {
    "Free-form (no cue)": "",
    "拳法 · Punches": "fist and punch techniques, fast hand strikes in continuous motion",
    "腿法 · Kicks": "kicking techniques, high leg strikes and spinning kicks",
    "综合套路 · Combination forms": "a continuous martial arts form, combination techniques flowing one into the next",
    "棍法 · Staff": "staff technique, spinning and striking with a long staff",
}
DEFAULT_TECHNIQUE = "Free-form (no cue)"


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
    frames = max(1, round(float(seconds) * FPS))
    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
        frames += 1
    return frames


def compose_prompt(prompt: str, technique: str = DEFAULT_TECHNIQUE) -> str:
    """The prompt the model is actually conditioned on: the request plus the technique-family cue."""
    prompt = (prompt or "").strip()
    cue = TECHNIQUES.get(technique or DEFAULT_TECHNIQUE, "")
    if not cue:
        return prompt
    return f"{prompt.rstrip('.,;')}, {cue}"


def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
    """Let the pipeline generate below its 5 s floor — the LoRA's own clips are 3.75 s."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

    MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))


# --- Global state ---
PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
LORA_STATUS: str | None = None


def status() -> str:
    if LOAD_ERROR:
        return LOAD_ERROR
    if PIPE is None:
        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 77.3 GB). Watch the Space logs."
    return (
        f"Ready · transformer + VAEs **bfloat16** · "
        f"placement `{PLACEMENT}` · "
        f"attention `{ATTENTION}` · "
        f"{LORA_STATUS or 'no LoRA'} · "
        f"loaded in {LOADED_IN:.0f}s · "
        f"conditioner `{CONDITIONER_SPACE}`"
    )


def load_models() -> str | None:
    """Load the denoising half at startup."""
    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, LORA_STATUS

    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR

    started = time.time()
    try:
        import torch
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3GeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()

        print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
        pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)

        # Fold the wushu action LoRA (+ the Turbo LoRA) into the bf16 weights.
        import h3_lora

        LORA_STATUS = h3_lora.apply_lora(pipe.transformer)
        if LORA_STATUS:
            print(f"[gen] {LORA_STATUS}", flush=True)

        pipe.transformer.set_attention_backend(ATTENTION)

        if PLACEMENT == "pack":
            pipe.transformer.to("cuda")

        PIPE = pipe
        MANAGER = manager
        LOADED_IN = time.time() - started
        print(f"[gen] ready in {LOADED_IN:.0f}s", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = (
            f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: "
            f"`{type(error).__name__}: {error}`"
        )
    return LOAD_ERROR


@cache
def conditioner():
    """The other half, over the gradio API."""
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def encode_remote(prompt, canvas, num_frames):
    """`/encode` on the conditioner Space — text only, this LoRA is text-to-video."""
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        image_path=None,
        last_image_path=None,
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=False,
        api_name="/encode",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


# --- GPU duration estimation ---
# Fitted on this Space's own ZeroGPU pool (RTX Pro 6000, 4 warm measurements of `_generate` wall time).
# `rows` is the packed latent sequence length, so per-step cost is linear + quadratic (attention) in it:
#
#     rows    steps  measured        rows    steps  measured
#     13770     4      23 s          32560     4      39 s
#     13770     8      33 s          32560     8    73.5 s
#
# -> per-step 2.5 s @ 13770 and 8.6 s @ 32560, which solves to the two coefficients below (+4% headroom).
# The residual fixed cost (audio/video decode, VAE placement) measured 4.5-13 s; 14 s covers it.
_DUR_B = 1.25e-4
_DUR_C = 4.6e-9
_DECODE_FIXED = 14
_PAD = 8
# The very first GPU call of the process also pays VAE placement + cuDNN attention autotune. That one-off is
# both large and very noisy -- the same request that takes 23-27 s warm measured 55 s, 63 s and 78 s cold on
# three different boots -- so it gets a deliberately wide allowance. It is charged once per process, so
# over-booking it costs almost nothing in aggregate while an under-book aborts somebody's first request.
_COLD_START = 70
_WARMED_UP = False


def get_duration(prompt_embeds, text_token_tags, height, width, num_frames, steps, seed, *a, **k):
    height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)

    latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
    rows = latent_frames * (height // 32) * (width // 32)
    denoise = steps * (_DUR_B * rows + _DUR_C * rows**2)
    seconds = denoise + _DECODE_FIXED + _PAD
    if not _WARMED_UP:
        seconds += _COLD_START
    return int(math.ceil(seconds))


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, height, width, num_frames, steps, seed):
    """The only thing on GPU time: the packed-sequence denoise loop and the two decoders."""
    import torch

    if PLACEMENT == "lazy":
        PIPE.to("cuda")
    elif PLACEMENT == "pack":
        PIPE.vae.to("cuda")
        PIPE.audio_vae.to("cuda")

    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        image=None,
        last_image=None,
        height=int(height),
        width=int(width),
        num_frames=int(num_frames),
        num_inference_steps=int(steps),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")


def generate(
    prompt,
    technique=DEFAULT_TECHNIQUE,
    canvas=DEFAULT_CANVAS,
    duration=TRAINED_DURATION,
    steps=DEFAULT_STEPS,
    seed=DEFAULT_SEED,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a martial-arts action clip with the MiniMax-H3 wushu action LoRA.

    Parameters:
        prompt: What the fighter does, e.g. "a kung fu practitioner executing a spinning kick"
        technique: Technique family cue appended to the prompt (punches / kicks / forms / staff)
        canvas: Output resolution and aspect ratio
        duration: Clip length in seconds, snapped to the 17n+5 frames the video VAE decodes
        steps: Denoising steps (4 with the Turbo LoRA folded in)
        seed: Random seed for reproducibility
    """
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading. Please wait a moment and try again.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please describe the martial-arts action you want, e.g. 'a fighter throws a spinning kick'.")

    from diffusers.utils import encode_video

    full_prompt = compose_prompt(prompt, technique)
    num_frames = snap_frames(duration)

    progress(0.0, desc=f"Conditioning on {CONDITIONER_SPACE} ...")
    conditioned = time.time()
    prompt_embeds, text_token_tags, metadata, plan = encode_remote(full_prompt, canvas, num_frames)
    condition_seconds = time.time() - conditioned
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))

    progress(0.1, desc=f"Denoising {int(steps)} steps at {width}x{height}, {num_frames} frames ...")
    started = time.time()
    frames, audio, sampling_rate = _generate(
        prompt_embeds, text_token_tags, height, width, num_frames, steps, seed
    )
    generate_seconds = time.time() - started

    # Placement and kernel autotune are paid once; later requests book the (much smaller) warm estimate.
    global _WARMED_UP
    _WARMED_UP = True

    directory = os.path.join(tempfile.gettempdir(), "h3-outputs")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"h3-wushu-{int(time.time() * 1000)}.mp4")
    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)

    report = (
        f"**Prompt sent to the model:** {full_prompt}\n\n"
        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.2f} s), {int(steps)} steps · "
        f"seed {int(seed)} · conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens) · "
        f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / max(int(steps), 1):.1f} s/step)"
    )
    print(f"[gen] {report}", flush=True)
    return path, report


# --- Load models at startup ---
load_models()

INTRO = """# 武打动作 · MiniMax-H3 Wushu Action LoRA

<div>
  <a href="https://huggingface.co/Jojocodex/minimax-h3-wushu-action-lora" target="_blank" rel="noopener"><strong>[ LoRA ]</strong></a> &nbsp;
  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank" rel="noopener"><strong>[ base model ]</strong></a> &nbsp;
  <a href="https://huggingface.co/Comfy-Org/MiniMax-H3" target="_blank" rel="noopener"><strong>[ ComfyUI weights ]</strong></a>
</div>

Generate short **martial-arts / kung-fu action** clips — punches, kicks, forms, staff work — with
[MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) and the
[Wushu Action LoRA](https://huggingface.co/Jojocodex/minimax-h3-wushu-action-lora), trained on 455 curated wushu
clips at 90 frames / 24 fps. The video comes back with H3's native synchronized soundtrack.

The LoRA has **no trigger token** — it activates on the action description itself. Describe the strike, then pick a
technique family to steer it toward the wording its captions used. Generation takes roughly a minute.
"""

CSS = """
.main.fillable {max-width: 1250px !important}
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="MiniMax-H3 Wushu Action LoRA") as demo:
    gr.Markdown(INTRO)

    with gr.Row(equal_height=True):
        with gr.Column():
            prompt = gr.Textbox(
                label="Action",
                lines=3,
                placeholder="e.g. 'a kung fu practitioner executing a spinning kick on a temple courtyard'",
                value="a martial artist performing punches and kicks in fast combat",
            )
            technique = gr.Radio(
                label="Technique family",
                info="Appended to the prompt in the wording the LoRA's captions used.",
                choices=list(TECHNIQUES),
                value=DEFAULT_TECHNIQUE,
            )
            run = gr.Button("Generate", variant="primary", size="lg")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                duration = gr.Slider(
                    label="Duration (s)",
                    info="Snapped up to the next 17n+5 frames. The LoRA was trained at 90 frames (3.75 s).",
                    minimum=MIN_UI_DURATION,
                    maximum=MAX_UI_DURATION,
                    step=0.25,
                    value=TRAINED_DURATION,
                )
                steps = gr.Slider(
                    label="Steps",
                    info="4 is enough with the Turbo LoRA folded in — more steps mostly just cost GPU time.",
                    minimum=4,
                    maximum=12,
                    step=5,
                    value=DEFAULT_STEPS,
                )
                seed = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)

        with gr.Column():
            video = gr.Video(label="Video + soundtrack", autoplay=True)
            report = gr.Markdown()

    gr.Examples(
        examples=[
            ["a martial artist performing punches and kicks in fast combat", "Free-form (no cue)"],
            ["a kung fu practitioner executing a spinning kick", "腿法 · Kicks"],
            ["two fighters exchanging strikes in an intense fight", "Free-form (no cue)"],
            ["a martial artist performing a powerful roundhouse kick, with explosive force", "腿法 · Kicks"],
            ["a practitioner demonstrating a fast flurry of punches in continuous motion", "拳法 · Punches"],
            ["a fighter executing a spinning staff technique", "棍法 · Staff"],
            ["a wushu athlete running through a form in a courtyard at dawn", "综合套路 · Combination forms"],
        ],
        inputs=[prompt, technique],
        outputs=[video, report],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

    inputs = [prompt, technique, canvas, duration, steps, seed]
    run.click(generate, inputs, [video, report], api_name="generate")
    prompt.submit(generate, inputs, [video, report], api_name=False)


if __name__ == "__main__":
    demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS, max_threads=1000)