app.py

"""MiniMax-H3 `ref2va`, split deployment — the denoising half.

This Space holds the `transformer_ref` partition and the two autoencoders, unquantized bfloat16. Text encoding runs in
[`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner), which this one calls over the
gradio API for every request; `reference_encoder` stays here, next to the autoencoders it runs.
"""

from __future__ import annotations

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

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to
# `ComponentsManager.enable_auto_cpu_offload`. Startup placement is not an option here — see `load_models`.
PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.
# flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
# Bounds on what `get_duration` may reserve. The pool reserves whatever number it is given, so a flat ceiling for every
# request is what makes an account hit "too many ZeroGPU credits allocated to running tasks".
MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))

# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
# is rejected there and surfaces as a failure here.
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"
EXAMPLE_CANVAS, EXAMPLE_PORTRAIT_CANVAS = "1344x768 · 16:9 full", "768x1024 · 3:4 full"
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
# It is the *snapped* frame count the ceiling has to hold for: 15 s is 360 frames, which rounds up to 362, i.e.
# 15.083 s, and is refused. 14 is the last whole second that survives the snap.
MAX_UI_DURATION = 14
MIN_DURATION = 2
# A reference video shorter than 2 s gives the model almost no motion to read.
MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
# `MINIMAX_H3_MAX_REFERENCE_IMAGES`. The slots are built up front and revealed one at a time, because a demo asking
# for two subjects should not open with nine boxes.
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2

# How many LoRA slots the UI offers, and the range each strength slider covers.
LORA_SLOTS = 3
LORA_MIN_SCALE, LORA_MAX_SCALE = -2.0, 2.0

# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
# matmuls, quadratic for the attention, against the AoTI block package this Space runs.
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call; every request carries it, because
# nothing here knows whether the worker it lands on is cold.
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
# Reading one adapter off local disk and injecting it across the 33B transformer's linear layers.
LORA_ALLOWANCE = 12


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 lower_duration_floor(seconds: float = MIN_DURATION) -> None:
    """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the released checkpoint."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

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


def video_latent_frames(num_frames: int) -> int:
    """`17 * n + 5` frames become `5 * n + 2` video latents."""
    return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2


def target_rows(height: int, width: int, num_frames: int) -> int:
    """The generated rows of the packed sequence: video patched `(1, 2, 2)`, plus two audio rows per latent."""
    video = video_latent_frames(num_frames) * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS


def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
    """The rows the reference blocks add, from metadata alone — no decode.

    An image is resized to a 2048 pixel short edge and encoded as a single frame; a video is put on the canvas *its
    own* aspect ratio resolves to, truncated to the generated frame count and snapped **down** to a `17 * n + 5` the
    VAE encodes without padding; a soundtrack contributes two rows per 1/40 s.
    """
    from PIL import Image

    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size

    rows = 0
    for kind, path in references:
        if kind == "image":
            width, height = Image.open(path).size
            scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
            resolved = [
                max(CANVAS_MULTIPLE, round(edge * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
                for edge in (height, width)
            ]
            rows += (resolved[0] // CANVAS_MULTIPLE) * (resolved[1] // CANVAS_MULTIPLE)
            continue

        video_seconds, audio_seconds = probe(path)
        if kind == "video" and video_seconds is not None:
            import av

            with av.open(path) as container:
                stream = container.streams.video[0]
                source_height, source_width = stream.height, stream.width
            canvas_height, canvas_width = resolve_canvas_size(source_width, source_height, CANVAS_MULTIPLE)
            frames = min(round(video_seconds * FPS), num_frames)
            snapped = max(1, (frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK
            rows += (
                video_latent_frames(snapped)
                * (canvas_height // CANVAS_MULTIPLE)
                * (canvas_width // CANVAS_MULTIPLE)
            )
        if audio_seconds is not None:
            seconds = min(audio_seconds, num_frames / FPS)
            rows += round(seconds * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
    return rows


def get_duration(
    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras=(), **_
):
    """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
    tolerates the `gr.Progress` `spaces` injects."""
    sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
        height, width, num_frames
    )
    denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
    # The two reference encoders ahead of the loop, and the two decoders plus the mux after it. Both scale with what
    # they are handed rather than with the step count.
    encode = 5 + reference_rows(references, num_frames) * 1e-3
    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10 + LORA_ALLOWANCE * len(loras or ())
    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
    return duration


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None


def load_models() -> str | None:
    """Load the denoising half at startup, but *not* onto the card.

    `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, the two schedulers and
    `video_processor`, so `load_components` fetches exactly those subfolders — `text_encoder/` and the `transformer/`
    partition are never touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a
    bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.

    Nothing moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()` writes every
    startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts the 150 GB quota
    (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack).
    """
    global PIPE, MANAGER, LOAD_ERROR

    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 MiniMaxH3Ref2VAGeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3Ref2VAGeneratorBlocks()
        print(f"[ref2va] 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)

        # Both VAEs first, and explicitly. `set_attention_backend` also sets the registry's *global* backend, which
        # every processor that was not stamped falls through to, and the float32 audio VAE has no cuDNN kernel:
        # `RuntimeError: No available kernel. Aborting execution.` in its causal encoder attention, which only a
        # reference soundtrack ever reaches.
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer_ref.set_attention_backend(ATTENTION)

        # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
        # worker. Off unless `H3_AOTI=1`. It is the *same* package the `transformer/` partition runs — the two configs
        # are identical field for field and the compiled code carries no weights of either.
        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer_ref)

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)

        PIPE, MANAGER = pipe, manager
        print(f"[ref2va] ready in {time.time() - started:.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


def _arm_decode_hooks(pipe):
    """Make the offload hooks fire for the two VAEs.

    `enable_auto_cpu_offload` wraps `forward`, and the reference-encoder and decode blocks call `vae.encode/decode(...)`
    directly, so the hook never runs and the VAE is still on the host when the latents arrive on the card.
    """
    for name in ("vae", "audio_vae"):
        module = getattr(pipe, name)
        for method in ("encode", "decode"):
            inner = getattr(module, method)

            def armed(*args, _module=module, _inner=inner, **kwargs):
                hook = getattr(_module, "_hf_hook", None)
                if hook is not None:
                    hook.pre_forward(_module)
                return _inner(*args, **kwargs)

            setattr(module, method, armed)


# ----------------------------------------------------------------------------------------------------------------
# LoRA
# ----------------------------------------------------------------------------------------------------------------
# There is no `MiniMaxH3LoraLoaderMixin` in the diffusers integration, so adapters are attached at the *model* level,
# through the `PeftAdapterMixin` the transformer carries. That is the whole API this needs: `load_lora_adapter` for
# each file and one `set_adapters` call to give them their strengths. Here the model is `transformer_ref`, so the
# adapters have to be trained against the `transformer_ref/` partition — a `transformer/` adapter is a different
# partition and will not match.


def _hub_url_parts(url: str) -> tuple[str, str]:
    """Split a huggingface.co `blob`/`resolve` URL into its repo id and the file path inside it."""
    from urllib.parse import unquote, urlparse

    parts = unquote(urlparse(url).path).strip("/").split("/")
    if len(parts) < 5 or parts[2] not in ("resolve", "blob"):
        raise gr.Error(f"Не разпознавам този адрес като файл в Hugging Face: `{url}`")
    return "/".join(parts[:2]), "/".join(parts[4:])


def resolve_lora(reference: str) -> str:
    """Turn what the user typed into a local `.safetensors` path.

    Accepts a local path, a huggingface.co file URL, `owner/repo/path/to/file.safetensors`, or a bare `owner/repo`
    whose single `.safetensors` is then picked for them. Runs outside the GPU call, so the download costs no GPU time.
    """
    from huggingface_hub import hf_hub_download, list_repo_files

    reference = (reference or "").strip()
    if not reference:
        return ""
    if os.path.exists(reference):
        return reference
    if reference.startswith(("http://", "https://")):
        repo_id, filename = _hub_url_parts(reference)
        return hf_hub_download(repo_id, filename)

    parts = [part for part in reference.split("/") if part]
    if len(parts) > 2 and parts[-1].endswith(".safetensors"):
        return hf_hub_download("/".join(parts[:2]), "/".join(parts[2:]))
    if len(parts) != 2:
        raise gr.Error(
            f"`{reference}` не е нито съществуващ файл, нито `автор/хранилище`, нито адрес към Hugging Face."
        )

    candidates = [name for name in list_repo_files(reference) if name.endswith(".safetensors")]
    if not candidates:
        raise gr.Error(f"В `{reference}` няма `.safetensors` файл.")
    if len(candidates) > 1:
        preferred = [name for name in candidates if "lora" in name.lower()]
        if len(preferred) != 1:
            listed = ", ".join(f"`{name}`" for name in sorted(candidates)[:8])
            raise gr.Error(f"`{reference}` съдържа няколко файла. Напиши `{reference}/име.safetensors`. Има: {listed}")
        candidates = preferred
    return hf_hub_download(reference, candidates[0])


def _lora_prefix(state_dict) -> str | None:
    """The prefix `load_lora_adapter` has to strip before the keys match the transformer's own module names."""
    key = next(iter(state_dict))
    for prefix in ("model.diffusion_model", "diffusion_model", "transformer_ref", "transformer"):
        if key.startswith(f"{prefix}."):
            return prefix
    return None


def apply_loras(transformer, loras) -> list[str]:
    """Attach `loras` (local path, strength) to `transformer` and give each its strength, replacing whatever was on it.

    Every adapter already on the model is removed first, so a request is never affected by the one before it — which
    matters when a worker is reused rather than forked fresh.
    """
    import torch

    from safetensors.torch import load_file

    for name in list(getattr(transformer, "peft_config", None) or {}):
        transformer.delete_adapters(name)

    names, scales = [], []
    for index, (path, scale) in enumerate(loras):
        state_dict = load_file(path)
        name = f"lora{index}"
        transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
        names.append(name)
        scales.append(float(scale))

    if not names:
        return []

    # PEFT builds the new layers on its own default device/dtype; the base weights are the truth here, under either
    # placement mode (`offload` keeps them on the host and moves whole modules by hook).
    base = next(param for key, param in transformer.named_parameters() if ".lora_" not in key)
    with torch.no_grad():
        for key, param in transformer.named_parameters():
            if ".lora_" in key and (param.device != base.device or param.dtype != base.dtype):
                param.data = param.data.to(device=base.device, dtype=base.dtype)

    transformer.set_adapters(names, scales)
    return names


def collect_loras(lora_fields, progress) -> tuple[list[tuple[str, float]], list[str]]:
    """Resolve the UI's `reference, strength, reference, strength, ...` into `(local path, strength)` pairs.

    Resolved before the booking: a download that happens inside `@spaces.GPU` is billed as GPU time.
    """
    loras, labels = [], []
    for reference, scale in zip(lora_fields[::2], lora_fields[1::2]):
        reference = (reference or "").strip()
        if not reference or abs(float(scale)) < 1e-6:
            continue
        progress(0.0, desc=f"Fetching LoRA {reference} ...")
        loras.append((resolve_lora(reference), float(scale)))
        labels.append(f"{os.path.basename(reference)} @ {float(scale):g}")
    if loras and os.environ.get("H3_AOTI") == "1":
        raise gr.Error("LoRA не може да се приложи върху AoTI компилиран трансформър. Изключи `H3_AOTI`.")
    return loras, labels


@cache
def conditioner():
    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, so the
    conditioner's booking is billed to whoever asked for the video."""
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def probe(path: str) -> tuple[float | None, float | None]:
    """`(video seconds, audio seconds)` of a media file, either being `None` when the stream is absent."""
    import av

    def seconds(stream, container):
        if stream.duration is not None and stream.time_base is not None:
            return float(stream.duration * stream.time_base)
        return None if container.duration is None else container.duration / av.time_base

    with av.open(path) as container:
        video = seconds(container.streams.video[0], container) if container.streams.video else None
        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
    return video, audio


def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
    """The `(kind, path)` references of a request, **in the order the model reads them**.

    That order numbers the labels of MiniMax-H3's prompt presentation and advances the shared audio/video rotary clock,
    so the same references in a different order are a different request.
    """
    ordered = [("image", path) for path in image_paths if path]
    if audio_path:
        ordered.append(("audio", audio_path))
    if video_path:
        ordered.append(("video", video_path))
    return ordered


def build_references(references: list[tuple[str, str]]):
    """The `(kind, path)` references of a request as decoded reference dataclasses, in packed order. `from_file` brings
    the rates along: a video its own frame rate and soundtrack, a clip its sample rate."""
    from diffusers.modular_pipelines.minimax_h3 import (
        MiniMaxH3AudioReference,
        MiniMaxH3ImageReference,
        MiniMaxH3VideoReference,
    )

    classes = {"image": MiniMaxH3ImageReference, "video": MiniMaxH3VideoReference, "audio": MiniMaxH3AudioReference}
    return [classes[kind].from_file(path) for kind, path in references]


def audio_bearing(references: list[tuple[str, str]]) -> list[tuple[str, float]]:
    """The references that carry a waveform, and how long it is. A video reference brings its own soundtrack."""
    carried = []
    for kind, path in references:
        if kind == "image":
            continue
        _, audio_seconds = probe(path)
        if audio_seconds is not None:
            carried.append((kind, audio_seconds))
    return carried


def duration_controls(audio_path, video_path, match: bool):
    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out."""
    try:
        carried = audio_bearing(collect([], audio_path, video_path))
    except Exception:
        carried = []
    # Exactly one soundtrack, long enough to be a duration MiniMax-H3 generates; anything else is ambiguous or out of
    # range and the slider stays.
    derivable = len(carried) == 1 and MIN_DURATION <= snap_frames(carried[0][1]) / FPS <= MAX_REFERENCE_VIDEO
    return gr.update(visible=derivable), gr.update(visible=not (derivable and match))


def check(prompt: str, references: list[tuple[str, str]]) -> None:
    """The model's own rules, before anything is uploaded or a card is allocated."""
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt, references or not.")
    if not references:
        raise gr.Error("Add at least one reference — an image or a video for the model to condition on.")
    if {kind for kind, _ in references} == {"audio"}:
        raise gr.Error("An audio reference needs an image or a video alongside it; it cannot go on its own.")
    for kind, path in references:
        if kind != "video":
            continue
        video_seconds, _ = probe(path)
        if video_seconds is None:
            raise gr.Error("That reference video has no video stream. Drop it in the audio slot instead.")
        if not MIN_REFERENCE_VIDEO <= video_seconds <= MAX_REFERENCE_VIDEO:
            raise gr.Error(
                f"The reference video is {video_seconds:.1f} s. Use a clip between "
                f"{MIN_REFERENCE_VIDEO:g} and {MAX_REFERENCE_VIDEO:g} seconds."
            )


def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
    """`/encode_ref2va` on the conditioner Space: a safetensors file holding `prompt_embeds` + `text_token_tags`, with
    the resolved `height` / `width` / `num_frames` in its metadata, plus the plan.

    `canvas` is the label. `media` and `kinds` are parallel and ordered, and the references go over because `ref2va`'s
    presentation puts a vision block in front of the prompt for every image and every merged video frame pair.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        media=[handle_file(path) for _, path in references],
        kinds=",".join(kind for kind, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        api_name="/encode_ref2va",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


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

    References cross as paths and are decoded here; only the three generated outputs come back. A `@spaces.GPU`
    argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is 370 MB of expanded frames, and
    the full `PipelineState` still holds the packed latents and the rotary grid on the card.

    The adapters are attached here rather than in the caller: `spaces` runs this body in its own worker, so the
    transformer the request sees is the one that has to carry them.
    """
    import torch

    if PLACEMENT == "lazy":
        PIPE.to("cuda")

    apply_loras(PIPE.transformer_ref, loras or ())

    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        references=build_references(references),
        height=height,
        width=width,
        num_frames=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(
    # The first four are the columns `gr.Examples` varies, and they lead the signature for that reason: an example row
    # is applied to `inputs` positionally. Every parameter has a default, which is what lets a four-column row call
    # this at all, and `upsample` is last so a positional API client that predates it is unaffected.
    prompt,
    image_1=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    match=True,
    duration=5,
    steps=28,
    seed=42,
    upsample=False,
    *lora_fields,
    progress=gr.Progress(track_tqdm=True),
):
    """One request. The LoRA fields are last and default to empty, so a positional API client that predates them is
    unaffected. `lora_fields` arrives as `reference, strength, reference, strength, ...`."""
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")

    from diffusers.utils import encode_video

    images = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9]
    references = collect(images, audio_path, video_path)
    check(prompt, references)

    # `0` is "leave it to the references" over the wire, which MiniMax-H3 accepts when exactly one of them carries a
    # soundtrack. The conditioner resolves it either way and this Space pins whatever comes back.
    derivable = len(audio_bearing(references)) == 1
    requested = 0 if (match and derivable) else snap_frames(duration)

    loras, lora_labels = collect_loras(lora_fields, progress)

    progress(0.0, desc="Upsampling the prompt ..." if upsample else "Reading the prompt and references ...")
    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            prompt, references, canvas, requested, rewrite_prompt=upsample
        )
    except gr.Error:
        raise
    except Exception as error:
        # gradio only puts the exception *type* on the wire, so the useful half of a conditioner-side failure is in
        # that Space's logs.
        traceback.print_exc()
        raise gr.Error(
            f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
            "Its logs carry the full traceback."
        ) from error
    condition_seconds = time.time() - conditioned
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
    refined = plan.get("refined_prompt") or ""

    progress(0.1, desc=f"Generating {num_frames / FPS:.1f} s at {width}x{height} ...")
    started = time.time()
    frames, audio, sampling_rate = _generate(
        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras
    )
    generate_seconds = time.time() - started

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

    print(
        f"[ref2va] {[kind for kind, _ in references]} · `{width}x{height}`, {num_frames} frames "
        f"({num_frames / FPS:.3f} s), {int(steps)} steps · conditioner {condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens{', upsampled' if refined else ''}) · "
        f"denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"
        f"{' · LoRA ' + ', '.join(lora_labels) if lora_labels else ''}",
        flush=True,
    )
    return path, refined, gr.update(visible=bool(refined))


def _fill_lora_slots(files, *current):
    """Drop `.safetensors` files on the uploader and their paths land in the first free slots, so a local adapter
    needs no typing at all."""
    slots = list(current)
    for path in files or []:
        for index, value in enumerate(slots):
            if not (value or "").strip():
                slots[index] = path
                break
    return [gr.update(value=value) for value in slots]


load_models()

INTRO = """# MiniMax-H3 Reference

<div align="center">
  <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank" rel="noopener"><strong>[ model ]</strong></a> &nbsp;
  <a href="https://www.minimax.io/blog/minimax-h3" target="_blank" rel="noopener"><strong>[ blog ]</strong></a> &nbsp;
  <a href="https://huggingface.co/spaces/multimodalart/minimax-h3" target="_blank" rel="noopener"><strong>[ text / image to video ]</strong></a>
</div>

**MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a
fully synchronized soundtrack (ambience, foley, speech). Bring your own subject, voice or camera move as a
reference.
"""

LORA_HELP = """Each slot takes a Hugging Face repo (`owner/repo`), a file inside one
(`owner/repo/name.safetensors`), a file URL, or a local path — or just drop the files below. A strength of `0`
switches a slot off without clearing it. Adapters have to be trained against the `transformer_ref/` partition.
"""

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

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

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="The character walks through a neon-lit street in the rain, humming to themselves",
            )
            upsample = gr.Checkbox(label="Upsample prompt", value=False)
            # One tab per modality, in the order the model reads them. A reference left in a tab that is not the open
            # one is still part of the request.
            with gr.Tabs():
                with gr.Tab("Images"):
                    # One `gr.Row`, so gradio splits the width evenly and wraps at `min_width` rather than leaving a
                    # hole where a hidden slot used to be.
                    with gr.Row():
                        images = [
                            gr.Image(
                                label="Subject, style or scene",
                                type="filepath",
                                min_width=180,
                                # Fixed, so a row that wraps to a single slot stays the size of a full one.
                                height=210,
                                visible=index < OPEN_IMAGE_SLOTS,
                            )
                            for index in range(MAX_IMAGE_SLOTS)
                        ]
                    add_image = gr.Button("+ Add another image", size="sm", variant="secondary")
                with gr.Tab("Audio"):
                    audio = gr.Audio(label="A voice or a piece of music", type="filepath")
                with gr.Tab("Video"):
                    video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")
            run = gr.Button("Generate", variant="primary")

            with gr.Accordion("LoRA", open=False):
                gr.Markdown(LORA_HELP)
                lora_references, lora_scales = [], []
                for slot in range(LORA_SLOTS):
                    with gr.Row():
                        lora_references.append(
                            gr.Textbox(label=f"LoRA {slot + 1}", placeholder="owner/repo", scale=3)
                        )
                        lora_scales.append(
                            gr.Slider(
                                label="Strength",
                                minimum=LORA_MIN_SCALE,
                                maximum=LORA_MAX_SCALE,
                                step=0.05,
                                value=1.0,
                                scale=2,
                            )
                        )
                lora_upload = gr.File(
                    label="Drop .safetensors here to fill the slots",
                    file_count="multiple",
                    file_types=[".safetensors"],
                    type="filepath",
                )

            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                match = gr.Checkbox(label="Match the reference soundtrack", value=True, visible=False)
                duration = gr.Slider(
                    label="Duration (s)", minimum=MIN_DURATION, maximum=MAX_UI_DURATION, step=1, value=5
                )
                steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)
                seed = gr.Number(label="Seed", value=42, precision=0)

        with gr.Column():
            result = gr.Video(label="Video + soundtrack")
            # An output, so it can be revealed only for a request that asked for a rewrite.
            with gr.Accordion("Upsampled prompt", open=False, visible=False) as upsampled_panel:
                upsampled = gr.Textbox(show_label=False, lines=8, interactive=False)

    open_slots = gr.State(OPEN_IMAGE_SLOTS)

    def reveal_image_slot(open_count):
        open_count = min(open_count + 1, MAX_IMAGE_SLOTS)
        return [
            open_count,
            *[gr.update(visible=index < open_count) for index in range(MAX_IMAGE_SLOTS)],
            gr.update(visible=open_count < MAX_IMAGE_SLOTS),
        ]

    add_image.click(reveal_image_slot, open_slots, [open_slots, *images, add_image], api_name=False)

    for control in (audio, video, match):
        control.change(
            duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
        )

    # `reference, strength, reference, strength, ...`, which is how `generate` unpacks them.
    lora_inputs = [field for pair in zip(lora_references, lora_scales) for field in pair]
    lora_upload.upload(_fill_lora_slots, [lora_upload, *lora_references], lora_references, api_name=False)

    # Same order as `generate`'s signature: the exampled five first, then the remaining image slots, then the LoRA
    # fields the `*lora_fields` tail collects.
    request = [
        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample, *lora_inputs
    ]

    gr.Examples(
        examples=[
            [
                "The character walks through a neon-lit street in the rain, humming to themselves",
                "examples/subject.png",
                None,
                None,
                EXAMPLE_CANVAS,
            ],
            [
                "The character speaks to camera in a quiet room, lips matching every word",
                "examples/subject.png",
                "examples/voice.wav",
                None,
                EXAMPLE_PORTRAIT_CANVAS,
            ],
            [
                "The character moves with the same camera push, down a rainy alley at night",
                "examples/subject.png",
                None,
                "examples/motion.mp4",
                EXAMPLE_CANVAS,
            ],
        ],
        inputs=[prompt, images[0], audio, video, canvas],
        outputs=[result, upsampled, upsampled_panel],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

    run.click(generate, request, [result, upsampled, upsampled_panel], api_name="generate")


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