sada9

Running on Zero

sada9 / avatar_tab.py
Transfer Bot
Moved to Hugging Face automatically
c3da9f3 10 days ago
40.3 kB
# ==========================================================================
# FILE: avatar_tab.py  — تبِ «🗣️ آواتار سخنگو» داخلِ اسپیسِ ترکیبی
# --------------------------------------------------------------------------
# همان اسپیسِ آواتارِ قبلی (MiniMax-H3 «fl2va» با صدای داده‌شده: یک عکس + صدای
# شما ← ویدیوی سخنگو، روی GPU)، با همه‌ی قابلیت‌هایش و همان کد:
#   • همان مدل (multimodalart/MiniMax-H3-Pruned)، همان LoRAی توربو (۷ گام)،
#     همان کاندیشنرِ راه‌دور، همان فیلترِ محتوای ncii_guard، همان بوم‌ها،
#     همان برآوردِ زمانِ GPU (get_duration)، همان «Upsample prompt»، همان نمونه.
#   • مسیرِ API همان قبلی است: /generate
#        ورودی: [prompt, image, audio, canvas, duration, steps, seed, upsample]
#        خروجی: [video, report, upsampled_prompt, upsampled_panel]
#
# فقط دو تفاوت (هر دو طبقِ درخواست / لازمه‌ی کنارِ هم بودن با آهنگ):
#
# ۱) «روشِ فایلِ .pt» (رندرِ مرحله‌ایِ ادامه‌پذیر، /animate_a2v، /forget_a2v،
#    دکمه‌ی تخمین و h3_a2v_resume.py) حذف شد — دقیقاً مثلِ حالتِ H3_RESUME=0
#    اسپیسِ قبلی: هر ویدیو یک مرحله‌ی GPU است و حداکثر ۱۴ ثانیه. (کیفیتِ پایین
#    در همان یک مرحله تمام می‌شود و فایلِ وضعیت لازم ندارد.)
#
# ۲) جای‌گذاریِ وزن‌ها روی GPU «تنبل» (H3_PLACEMENT=lazy) است، نه «pack»:
#    ZeroGPU هر وزنی را که موقعِ بالا آمدن روی "cuda" گذاشته شود در یک بسته‌ی
#    مشترک روی دیسک می‌گذارد و در «هر» کارِ GPU (آهنگ یا آواتار) همه‌ی آن بسته
#    را روی کارت می‌برد. ترنسفورمرِ آواتار ۵۱ گیگابایت است؛ اگر در بسته می‌رفت:
#      • هر ساختِ آهنگ (کارتِ large، ۴۸ گیگ) با ۱۶ + ۵۱ گیگ وزن پُر/خراب می‌شد،
#      • بسته‌ی ۶۷ گیگی از آستانه‌ی xlarge (۴۲ گیگ) بالا می‌رفت و زمانِ رزروِ آهنگ
#        را هم عوض می‌کرد،
#      • دیسک (وزن‌های دانلودی + نسخه‌ی دومِ بسته) به سقفِ ۱۵۰ گیگِ ZeroGPU می‌رسید.
#    پس آواتار وزن‌هایش را خودش، داخلِ همان مرحله‌ی GPUِ خودش (کارتِ xlarge، مثلِ
#    قبل) روی کارت می‌برد و ساختِ آهنگ دقیقاً مثلِ قبل فقط وزن‌های خودش را دارد.
#    کارگرِ GPUِ آواتار هم وزن‌های آهنگ را روی کارتش نمی‌برد (هوکِ _move پایین)،
#    پس کلِ کارت مثلِ قبل در اختیارِ آواتار است. وزن‌های آواتار روی CPU از فایل‌های
#    safetensors نگاشت (mmap) می‌شوند؛ یعنی کشِ فایل‌اند نه حافظه‌ی ثابت، و سیستم هر
#    وقت رم لازم داشت می‌تواند آزادشان کند.
#
# ترتیبِ بالا آمدن (app.py): موتورِ صدا و فیلترِ محتوا (هر دو پروسه‌ی جدا، فقط
# CPU) همان اول؛ بعد مدل‌های آهنگ؛ بعد بسته‌بندیِ ZeroGPU؛ و بعد همین مدلِ آواتار
# (start_after_zerogpu_pack). یعنی بیشترین مصرفِ رم هنگامِ بالا آمدن پایین می‌ماند.
# ==========================================================================
"""MiniMax-H3 `fl2va` with a **given** soundtrack — a first frame plus your audio, split deployment, denoising half.
Integrated as one tab of the combined Space (music on GPU, voice cloning on CPU, this avatar on GPU). The code is the
avatar Space's own, minus the resumable `.pt` path (`H3_RESUME=0` behaviour: one GPU call per video, up to 14 s) and
with lazy placement, so its 51 GB transformer never enters the ZeroGPU pack the music tab's workers load.
"""

from __future__ import annotations

import contextvars
import os
import tempfile
import threading
import time
import traceback
from functools import cache

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the load can happen at startup
# rather than on GPU time. (app.py has already imported it first; importing it again is a no-op.)
import spaces
import gradio as gr

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "multimodalart/MiniMax-H3-Pruned")
# The pruned DiT's AoTI package is keyed apart from the released one: its AdaLN projections read 8 inputs instead of
# 2688, and `LazyAOTIModel` binds weights by fully qualified name, so loading the released package against these
# blocks is a SIGSEGV rather than an error. Set before `h3_aoti` is imported anywhere.
os.environ.setdefault("H3_WIDTH", "bf16-pruned")
os.environ.setdefault("H3_AOTI_KEY", "bf16-pruned/torch2.11/sm120/dynamic")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `pack` places the transformer at startup, `lazy` moves everything on the first GPU call.
# ALPHA: `lazy` by default in the combined Space — see the header. (`H3_PLACEMENT=pack` is still honoured, but then
# the music tab's ZeroGPU workers would load these 51 GB too.)
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.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# The turbo LoRA this Space runs. `larryvrh/MiniMax-H3-Turbo-Lora` v4 was the only leg of the ten-leg a2v sweep that
# beat its own step-matched control rather than merely matching it (mouth-motion proxy 0.4308 vs 0.2623, sharpness
# 252.5 vs 230.3), while still ~3.7x faster than no LoRA at all.
#
# `num_inference_steps` counts sigma grid points **including the terminal one, which has no model evaluation**, so a
# run makes `steps - 1` forward passes (NFE). 6 NFE is 7 steps. That is the convention throughout MiniMax-H3 and is
# not a typo here.
#
# What a correct attach looks like: 363 modules — 51 AdaLN + 312 attention/feed-forward — and 51 AdaLN projection
# offsets. The offsets are there because this LoRA is *released*-trained: on the pruned checkpoint served here its
# AdaLN factors go through the shipped full→pruned projection
# (`MiniMaxH3PrunedTransformer3DModel.project_adaln_lora`), whose constant term cannot live in the projection itself.
# A file that only partly claims the model loads without complaint and costs quality silently, so both counts are
# checked rather than trusted.
TURBO_REPO = "larryvrh/MiniMax-H3-Turbo-Lora"
TURBO_WEIGHT = "minimax_h3_turbo_v4_step600_ema.safetensors"
TURBO_ADAPTER = "larryvrh_v4"
TURBO_MODULES, TURBO_OFFSETS = 363, 51
TURBO_STEPS, BASE_STEPS = 7, 28
# `H3_TURBO=0` runs the base model instead, at the base step count. Env only — a debugging escape hatch, not a
# user-facing choice.
TURBO = os.environ.get("H3_TURBO", "1") == "1"

# The audio rows are always *locked*: the given soundtrack is written into them clean and held there, and the model
# is told those rows are already finished. `h3_a2v_blocks.py` still implements `blended` and `off` — the sweep code
# imports it — but nothing here reaches them.
AUDIO_CONDITION_MODE = "locked"

# Must stay identical to the conditioner's table: the *label* goes over the wire.
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),
    "1024x1024 · 1:1 max": (1024, 1024),
    # 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, 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 and is
# refused. 14 s is therefore the longest a *single* MiniMax-H3 request (one segment) can be.
MIN_UI_DURATION, MAX_UI_DURATION = 2, 14

EXAMPLE_IMAGE = os.path.join("examples", "avatar", "subject.png")
EXAMPLE_AUDIO = os.path.join("examples", "avatar", "voice.wav")


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_UI_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))


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
TURBO_STATUS = "off (`H3_TURBO=0`)" if not TURBO else "not loaded"
TURBO_ON = False


def default_steps() -> int:
    """The step count the model that actually loaded wants: the turbo LoRA's 6 NFE, or the base schedule."""
    return TURBO_STEPS if TURBO_ON else BASE_STEPS


def _expected_default_steps() -> int:
    """What `default_steps()` will say once the model is up (the tab is built before the model loads)."""
    if PIPE is not None or LOAD_ERROR is not None:
        return default_steps()
    return TURBO_STEPS if TURBO else BASE_STEPS


def check_prompt(prompt: str) -> None:
    """The NCII guard: every request here carries a keyframe, i.e. is animation on top of a real photo. It runs
    before the conditioner call and the denoise booking, so a refused prompt costs no GPU time on either half."""
    import ncii_guard

    flag = ncii_guard.classify(prompt)
    if flag["label"] == "ncii":
        print(f"[guard] prompt refused (ncii {flag['score']:.2f}): {prompt!r}", flush=True)
        raise gr.Error("This prompt was flagged by a content filter and wasn't run.")


def start_guard_async() -> None:
    """ALPHA: the original Space started the guard at import and blocked on it. Here it starts in the background
    with the rest of the startup, so it never delays (or, if it cannot load, takes down) the music and voice tabs.
    `classify` waits for it and revives it if it died — a request is never run without the guard's answer."""

    def _run():
        try:
            import ncii_guard

            ncii_guard.start()
            print("[guard] ready", flush=True)
        except Exception as error:
            print(f"[guard] could not start yet ({type(error).__name__}: {error}); the first avatar request retries",
                  flush=True)

    threading.Thread(target=_run, daemon=True, name="ncii-guard-start").start()


def detach_turbo(transformer) -> None:
    """Delete the adapter *and* the AdaLN constant terms that ride alongside it.
    `delete_adapters` is PEFT's, and PEFT knows nothing about the pruned checkpoint's `_lora_adaln_offsets`: those are
    plain non-persistent buffers this checkpoint's own loader registers for a released-trained LoRA. Dropping them
    here keeps "the adapter was removed" true of the whole adapter, including the half of it that does not live in a
    `lora_A`/`lora_B` pair.
    """
    if TURBO_ADAPTER in (getattr(transformer, "peft_config", None) or {}):
        transformer.delete_adapters(TURBO_ADAPTER)
    for _, module in transformer.named_modules():
        registry = getattr(module, "_lora_adaln_offsets", None)
        if isinstance(registry, dict) and TURBO_ADAPTER in registry:
            module._buffers.pop(registry.pop(TURBO_ADAPTER), None)


def attach_turbo(pipe) -> str:
    """Inject the turbo LoRA, and count what actually attached.
    Runs in the parent process, before the transformer goes to the card, so every forked ZeroGPU worker inherits it.
    """
    started = time.time()
    pipe.load_lora_weights(TURBO_REPO, weight_name=TURBO_WEIGHT, adapter_name=TURBO_ADAPTER)
    modules, offsets, scalings = [], [], set()
    for path, module in pipe.transformer.named_modules():
        scaling = getattr(module, "scaling", None)
        if isinstance(scaling, dict) and TURBO_ADAPTER in scaling:
            modules.append(path)
            scalings.add(round(float(scaling[TURBO_ADAPTER]), 6))
        if TURBO_ADAPTER in getattr(module, "_lora_adaln_offsets", {}):
            offsets.append(path)
    adaln = sum(path.endswith(("adaln_proj.linear", "norm_out.linear")) for path in modules)
    census = (
        f"{len(modules)} modules ({adaln} AdaLN + {len(modules) - adaln} attention/FF), "
        f"{len(offsets)} AdaLN projection offsets, scale {sorted(scalings)}"
    )
    print(f"[turbo] {TURBO_REPO}/{TURBO_WEIGHT}: {census}, in {time.time() - started:.1f}s", flush=True)
    if len(modules) != TURBO_MODULES or len(offsets) != TURBO_OFFSETS:
        # Loud rather than quiet: the Space keeps serving, without the adapter and at the base step count.
        detach_turbo(pipe.transformer)
        print(
            f"[turbo] expected {TURBO_MODULES} modules and {TURBO_OFFSETS} AdaLN offsets, got {len(modules)} and "
            f"{len(offsets)} — adapter deleted",
            flush=True,
        )
        return (
            f"**under-attached** ({len(modules)} of {TURBO_MODULES} modules, {len(offsets)} of {TURBO_OFFSETS} "
            "AdaLN offsets), removed — running the base model"
        )
    return f"`{TURBO_REPO}` · {census}"


def load_models() -> str | None:
    """Load the denoising half (once).
    `MiniMaxH3A2VGeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and the processors,
    so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` 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, and this Space *encodes* with it as well.
    """
    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, TURBO_STATUS, TURBO_ON

    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_a2v_blocks import MiniMaxH3A2VGeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3A2VGeneratorBlocks()
        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")
        # The pruned DiT is served as remote code (`transformer/modeling_minimax_h3_pruned.py`, reached through the
        # `AutoModel` type hint in `modular_model_index.json`). `load_components` forwards `trust_remote_code` only to
        # components that live in the pipeline's own repo, so the two VAEs and the schedulers — which
        # `modular_model_index.json` still points at `MiniMaxAI/MiniMax-H3` — never see it.
        pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)
        pipe.transformer.set_attention_backend(ATTENTION)

        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer)

        if TURBO:
            try:
                TURBO_STATUS = attach_turbo(pipe)
                TURBO_ON = TURBO_STATUS.startswith("`")
            except Exception as error:
                traceback.print_exc()
                TURBO_STATUS = f"**failed to load**: `{type(error).__name__}: {error}` — running the base model"
                TURBO_ON = False

        if PLACEMENT == "pack":
            # Scoped to the transformer: `spaces` packs every startup-resident CUDA tensor into a second on-disk copy,
            # and packing all 51.3 GB busts the 150 GB storage quota. The ~10 GB of fp32 VAEs move on the first call.
            pipe.transformer.to("cuda")

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


@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 encode_remote(prompt, image_path, canvas, num_frames, rewrite_prompt=False):
    """`/encode` 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.
    With `rewrite_prompt`, the conditioner first rewrites the prompt into MiniMax-H3's trained format on its own
    Qwen3-VL — looking at the keyframe, which is why the rewrite belongs there and not here — encodes *that*, and
    returns it under the plan's `refined_prompt`. It is a second `@spaces.GPU` booking on the conditioner, ahead of
    the encode's own; both are on that Space's clock and neither touches this one's, which is only ever the denoise.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        # No keyframe is a `t2va` request on the far side too: the conditioner's own auto blocks take the text-only
        # branch and the canvas comes from the label alone, since nothing has an aspect ratio to follow.
        image_path=handle_file(image_path) if image_path else None,
        last_image_path=None,
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        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


def load_waveform(audio_path):
    """The uploaded soundtrack as MiniMax-H3's stereo float32 at the audio VAE's 32 kHz."""
    from diffusers.modular_pipelines.minimax_h3.references import MiniMaxH3AudioReference

    from h3_a2v_blocks import normalize_waveform

    reference = MiniMaxH3AudioReference.from_file(audio_path)
    return normalize_waveform(reference.audio, reference.sample_rate, 32000), 32000


# ----------------------------------------------------------------------------------------------------------
# ZeroGPU and the music tab, side by side (ALPHA)
# ----------------------------------------------------------------------------------------------------------
# A ZeroGPU worker, on its first task, moves *every* packed tensor of the Space onto its card
# (`spaces.zero.torch.patching._move`). The pack holds the music models (~16 GB); the avatar worker never runs them,
# so for the avatar's own workers that move is skipped: `generate` marks its call with this context variable, the
# worker is forked from that very thread, and the hook below reads the mark in the child. Every other worker (the
# music tab's) moves exactly as before. If the hook cannot be installed (another `spaces` version), the music weights
# are dropped from the avatar's card right after the move instead (`_drop_foreign_gpu_tensors`).
_AVATAR_GPU_CALL: contextvars.ContextVar[bool] = contextvars.ContextVar("alpha_avatar_gpu_call", default=False)
_MOVE_HOOK = {"installed": False}
_FOREIGN_DROPPED = {"done": False}


def install_zerogpu_move_hook() -> bool:
    if PLACEMENT != "lazy" or _MOVE_HOOK["installed"]:
        return _MOVE_HOOK["installed"]
    try:
        from spaces.config import Config as SpacesConfig

        if not SpacesConfig.zero_gpu:
            return False
        from spaces.zero.torch import patching as zerogpu_patching
    except Exception as error:
        print(f"[avatar] ZeroGPU move hook not installed ({type(error).__name__}: {error})", flush=True)
        return False
    original = getattr(zerogpu_patching, "_move", None)
    if not callable(original):
        print("[avatar] ZeroGPU move hook not installed (spaces has no `_move`)", flush=True)
        return False

    def _move(*args, **kwargs):
        if _AVATAR_GPU_CALL.get():
            # The avatar's own worker: it places its own weights (lazy) and has no use for the music pack.
            return None
        return original(*args, **kwargs)

    _move._alpha_original = original
    zerogpu_patching._move = _move
    _MOVE_HOOK["installed"] = True
    print("[avatar] ZeroGPU: the avatar's GPU workers will not load the music weights", flush=True)
    return True


def _drop_foreign_gpu_tensors() -> None:
    """Inside the avatar's GPU worker only: if ZeroGPU still moved the music pack onto this card (no hook), free it.
    Those tensors belong to modules this worker never runs (a ZeroGPU worker serves one decorated function only), so
    replacing their storage frees the card for the avatar exactly as it had it before."""
    if PLACEMENT != "lazy" or _FOREIGN_DROPPED["done"]:
        return
    _FOREIGN_DROPPED["done"] = True
    try:
        import torch
        from spaces.zero.torch import patching as zerogpu_patching
    except Exception:
        return
    freed, fakes = 0, []
    try:
        for tensor_pack in list(getattr(zerogpu_patching, "tensor_packs", []) or []):
            for fake_list in (getattr(tensor_pack, "fakes", None) or {}).values():
                fakes.extend(fake_list)
        fakes.extend(list(getattr(zerogpu_patching, "cuda_aliases", {}).keys()))
    except Exception:
        return
    for fake in fakes:
        try:
            if fake.is_cuda and fake.numel():
                freed += fake.numel() * fake.element_size()
                fake.data = torch.empty(0, dtype=fake.dtype, device=fake.device)
        except Exception:
            pass
    if freed:
        torch.cuda.empty_cache()
        print(f"[avatar] freed {freed / 2**30:.1f} GiB of music weights from the avatar's card", flush=True)


# Seconds of GPU one request needs, from the packed video rows it is about to denoise: linear in the rows for the
# matmuls, quadratic for the attention. Fitted against the AoTI package the sibling Space runs, so a `H3_AOTI=0`
# deployment — this one, by default — pays the ~15 % eager penalty on top.
_DUR_B, _DUR_C, _EAGER = 1.1745e-4, 3.8396e-9, 1.2
# The two resident decoders and the mux, which scale with the output rather than with the step count.
_DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 15, 15, 960 * 544 * 124
# `pack` mode: only the ~10 GB of VAEs move on a cold worker.
_PLACEMENT_ALLOWANCE, _PAD = 12, 10


def _as_int(value, default):
    """ALPHA: a number from a request, or `default` when it came empty (None / "" / not a number). The logs showed
    requests reaching the GPU step with an empty seed (`int(None)` -> TypeError, and every such request failed after
    its conditioner booking was already spent); an empty seed or step count must never fail a video."""
    if value is None or isinstance(value, bool):
        return default
    try:
        return int(float(value))
    except (TypeError, ValueError):
        return default


def get_duration(prompt_embeds, text_token_tags, image, waveform, sample_rate, height, width, num_frames, steps, seed, *a, **k):
    # `steps` is the value in *this request*, not the turbo LoRA's 7: the slider stays adjustable and a hand-raised
    # count has to reserve the GPU it will actually use. Prompt upsampling costs nothing here — it happens on the
    # conditioner's clock, before this function is ever called.
    height, width, num_frames = int(height), int(width), int(num_frames)
    steps = _as_int(steps, default_steps())
    latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
    patches = (height // 32) * (width // 32)
    # The target video rows, plus the keyframe's one latent frame of condition rows when there is a keyframe. A
    # `t2va` request packs `[text | target audio | target video]` and has no condition block at all, so the row
    # count — and with it the quadratic attention term — must not assume one.
    rows = latent_frames * patches + (patches if image is not None else 0)
    denoise = steps * (_DUR_B * rows + _DUR_C * rows**2) * (1.0 if os.environ.get("H3_AOTI") == "1" else _EAGER)
    decode = _DECODE_BASE + _DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / _DEFAULT_CANVAS_PIXELS
    return max(60, int(denoise + decode) + _PLACEMENT_ALLOWANCE + _PAD)


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, image, waveform, sample_rate, height, width, num_frames, steps, seed):
    """The only thing on GPU time: the audio encode, the packed-sequence denoise loop and the two decoders.
    Only the frames come back — a `@spaces.GPU` return crosses a process boundary by pickling, and the full
    `PipelineState` still holds the packed latents, the rotary grid and the row indices on the card. The decoded
    audio rows are left behind too: with the rows given and locked, they are the upload round-tripped through the
    audio VAE, and it is the upload itself that goes onto the file.
    """
    import torch

    if seed is None or steps is None:
        print(f"[gen] ⚠️ the GPU step got an empty value (steps={steps!r}, seed={seed!r}); using the defaults",
              flush=True)
    steps, seed = _as_int(steps, default_steps()), _as_int(seed, 42)

    if PLACEMENT == "lazy":
        _drop_foreign_gpu_tensors()
        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=image,
        given_audio=waveform,
        given_audio_sample_rate=sample_rate,
        audio_condition_mode=AUDIO_CONDITION_MODE,
        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]


def generate(prompt, image_path, audio_path, canvas=DEFAULT_CANVAS, duration=6, steps=None, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True)):
    """One request, one GPU call (up to 14 s) — exactly the avatar Space's `/generate` with `H3_RESUME=0`."""
    if seed is None or seed == "" or steps == "":
        print(f"[gen] ⚠️ request without a seed/steps value (seed={seed!r}, steps={steps!r}); using the defaults",
              flush=True)
    # ALPHA: an empty seed / step count falls back to the defaults instead of failing the video (see `_as_int`).
    steps = _as_int(steps, default_steps())
    seed = _as_int(seed, 42)
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")
    if not audio_path:
        raise gr.Error("A soundtrack is required.")
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt, keyframe or not.")
    check_prompt(prompt)

    if snap_frames(duration) > snap_frames(MAX_UI_DURATION):
        raise gr.Error(f"Videos longer than {MAX_UI_DURATION} s are not supported: one request is one segment.")

    from PIL import Image, ImageOps

    from diffusers.utils import encode_video

    num_frames = snap_frames(duration)
    waveform, sample_rate = load_waveform(audio_path)

    progress(0.0, desc="Upsampling the prompt ..." if upsample else f"Conditioning on {CONDITIONER_SPACE} ...")
    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            prompt, image_path, canvas, num_frames, 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. Worth saying so: with a rewrite asked for there are two bookings over there that can
        # fail, not one.
        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"))
    # `None` whenever no rewrite happened — whether none was asked for, or one was asked for and failed.
    refined = plan.get("refined_prompt") or ""

    # The conditioning latents encoded here have to be of the image the conditioner looked at, which prepares it
    # exactly this way. `None` is a `t2va` request: the blocks branch on it and skip the keyframe path entirely.
    keyframe = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB") if image_path else None

    progress(0.1, desc=f"Denoising {steps} steps at {width}x{height}, {num_frames} frames ...")
    started = time.time()
    # ALPHA: mark this call so the ZeroGPU worker it forks skips the music pack (see `install_zerogpu_move_hook`).
    mark = _AVATAR_GPU_CALL.set(True)
    try:
        frames = _generate(
            prompt_embeds, text_token_tags, keyframe, waveform, sample_rate, height, width, num_frames, steps, seed
        )
    finally:
        _AVATAR_GPU_CALL.reset(mark)
    generate_seconds = time.time() - started

    # The soundtrack on the file. With the audio rows given, the model's own audio output *is* the given waveform
    # round-tripped through the audio VAE, so the file carries the upload itself, cut to the video's own length.
    import torch

    samples = int(round(num_frames / FPS * sample_rate))
    audio = torch.nn.functional.pad(waveform, (0, max(0, samples - waveform.shape[-1])))[:, :samples]
    rate = sample_rate

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

    report = (
        f"`{'fl2va' if image_path else 't2va'}` · "
        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {int(steps)} steps, seed {int(seed)} · "
        f"soundtrack on the file: your upload · "
        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
        f"{', upsampled' if refined else ''}) · "
        f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · "
        f"turbo LoRA {'on' if TURBO_ON else 'off'}"
    )
    print(f"[gen] {report}", flush=True)
    return path, report, refined, gr.update(visible=bool(refined))


def _cleanup_old_outputs(directory, max_age_seconds=6 * 3600):
    """ALPHA: the videos are copied into Gradio's cache when they are returned, so the ones older than a few hours
    are only disk the Space keeps paying for (the original kept its outputs on the `/data` add-on via the resume
    path's retention; there is no such path here any more)."""
    try:
        cutoff = time.time() - max_age_seconds
        for name in os.listdir(directory):
            item = os.path.join(directory, name)
            try:
                if os.path.isfile(item) and os.path.getmtime(item) < cutoff:
                    os.remove(item)
            except Exception:
                pass
    except Exception:
        pass


def audio_duration(audio_path, current):
    """The duration control follows the upload: MiniMax-H3 generates a whole `17 * n + 5` frames, so the slider lands
    on the second that covers the clip, inside the model's own range."""
    if not audio_path:
        return gr.update()
    try:
        from diffusers.modular_pipelines.minimax_h3.references import _decode_audio_file

        waveform, rate = _decode_audio_file(audio_path)
        seconds = waveform.shape[-1] / rate
    except Exception as error:
        print(f"[ui] could not read {audio_path}: {error}", flush=True)
        return gr.update()
    return gr.update(value=min(MAX_UI_DURATION, max(MIN_UI_DURATION, round(seconds + 0.49))))


def fit_keyframe(image_path, current_canvas):
    """Cover-crop an uploaded first frame to the closest supported aspect ratio and select that ratio's smallest
    (fastest) canvas, unless the user already picked a matching ratio."""
    if not image_path:
        return gr.update(), gr.update()
    from PIL import Image as _Image

    img = _Image.open(image_path)
    aspect = img.width / img.height
    fastest = {}
    for label, (h, w) in CANVASES.items():
        r = w / h
        if r not in fastest or w * h < fastest[r][1][0] * fastest[r][1][1]:
            fastest[r] = (label, (h, w))
    ratio = min(fastest, key=lambda r: abs(r - aspect))
    label, (h, w) = fastest[ratio]

    cur_h, cur_w = CANVASES[current_canvas]
    if abs(cur_w / cur_h - aspect) <= abs(ratio - aspect):
        label, h, w = current_canvas, cur_h, cur_w

    target = w / h
    if abs(img.width / img.height - target) <= 1e-3:
        return gr.update(), gr.update(value=label)
    if img.width / img.height > target:
        new_w = int(img.height * target)
        left = (img.width - new_w) // 2
        img = img.crop((left, 0, left + new_w, img.height))
    else:
        new_h = int(img.width / target)
        top = (img.height - new_h) // 2
        img = img.crop((0, top, img.width, top + new_h))
    img.save(image_path)
    return gr.update(value=image_path), gr.update(value=label)


# ----------------------------------------------------------------------------------------------------------
# Startup (called from app.py)
# ----------------------------------------------------------------------------------------------------------
def start_after_zerogpu_pack() -> bool:
    """Load the model right after ZeroGPU has packed (and freed from RAM) the music weights, still before the server
    takes its first request — the lowest RAM peak that keeps "everything is up when the Space is up".
    Returns False when there is no ZeroGPU pack step to hook into (then app.py loads the model before launch)."""
    try:
        from spaces.config import Config as SpacesConfig

        if not SpacesConfig.zero_gpu:
            return False
        import spaces.zero.torch as zerogpu_torch
    except Exception:
        return False
    original_pack = getattr(zerogpu_torch, "pack", None)
    if not callable(original_pack):
        return False

    def pack(*args, **kwargs):
        total = original_pack(*args, **kwargs)
        print("[avatar] ZeroGPU packed the music weights; loading the avatar model now ...", flush=True)
        load_models()
        return total

    zerogpu_torch.pack = pack
    return True


INTRO = """# MiniMax-H3 · audio to video
<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> &nbsp;
  <a href="https://huggingface.co/spaces/multimodalart/minimax-h3-reference" target="_blank" rel="noopener"><strong>[ reference to video ]</strong></a>
</div>
**MiniMax-H3** denoises video and audio as one packed sequence, so a soundtrack can be handed to it instead of
generated. Bring your own audio, and optionally a first frame.
<div dir="rtl" style="text-align:right">
**آواتار سخنگو:** یک عکس (فریمِ اول) + صدای خودتان ← ویدیوی سخنگو، روی GPU. هر ویدیو در **یک مرحله** ساخته
می‌شود (حداکثر ۱۴ ثانیه). برای سهمیه‌ی کم: بوم کوچک، مدت کوتاه، «Upsample prompt» خاموش.
</div>
"""


def build_avatar_tab(demo):
    """The avatar Space's own page, inside the combined Space's third tab (without the resumable `.pt` section)."""
    install_zerogpu_move_hook()
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                image = gr.Image(label="First frame", type="filepath", height=260)
                audio = gr.Audio(label="Soundtrack", type="filepath")
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe the shot.")
            upsample = gr.Checkbox(label="Upsample prompt", value=False)
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                duration = gr.Slider(
                    label="Duration (s)",
                    minimum=MIN_UI_DURATION,
                    maximum=MAX_UI_DURATION,
                    step=1,
                    value=6,
                    info="Follows the soundtrack on upload. A longer video pads the audio with silence.",
                )
                steps = gr.Slider(
                    label="Steps",
                    minimum=3,
                    maximum=40,
                    step=1,
                    value=_expected_default_steps(),
                    info="Sigma grid points; one fewer forward pass than this. The turbo LoRA's card says 6–8 steps.",
                )
                seed = gr.Number(label="Seed", value=42, precision=0)

        with gr.Column():
            video = gr.Video(label="Video")
            # 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)
            report = gr.Markdown(elem_id="h3-report")

    image.upload(fit_keyframe, [image, canvas], [image, canvas])
    # `change` rather than `upload`, so picking the example updates the duration too.
    audio.change(audio_duration, [audio, duration], [duration])

    gr.Examples(examples=[[EXAMPLE_IMAGE, EXAMPLE_AUDIO]], inputs=[image, audio], label="Inputs")

    run.click(
        generate,
        [prompt, image, audio, canvas, duration, steps, seed, upsample],
        [video, report, upsampled, upsampled_panel],
        api_name="generate",
    )

    # The tab is built before the model loads (it loads right after ZeroGPU's pack step): once it is up, a fresh page
    # shows the step count the loaded model really wants (7 with the turbo LoRA, 28 without it) — as before.
    demo.load(lambda: gr.update(value=default_steps()), None, [steps], api_name=False, show_progress="hidden",
              queue=False)
    return {"run": run, "image": image, "audio": audio, "video": video}