"""MiniMax-H3 Talking Avatar — minimal single-purpose Space.

Only two inputs: an image (the character) and an audio file (the voice). One "Generate" button produces a
lip-synced talking-avatar video. Everything else (multi-image references, video references, canvas choice,
duration, steps, seed, prompt upsampling) is fixed to sane defaults internally — there is nothing else to
configure in the UI on purpose.

This keeps the split-deployment architecture of the original multimodalart/minimax-h3-reference Space: the
33B model is ~196 GiB in bf16, far past what a single ZeroGPU worker can hold, so text encoding (the 62 GiB
Qwen3-VL half) runs on a separate Space (`multimodalart/qwen3vl-conditioner`) over the Gradio API, and only
the denoising half (`transformer_ref` + the two autoencoders) loads here.
"""

from __future__ import annotations

import os
import random
import tempfile
import time
import traceback

import spaces  # noqa: F401  (must import before anything touches torch.cuda)
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")
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn")
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))

# All canvases MiniMax-H3 was trained on. (height, width) per label — the model can only output one of these,
# never an arbitrary resolution, so we pick whichever one's aspect ratio is closest to the uploaded image's own.
CANVASES = {
    "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),
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
# Prefer the "full" quality tier when it ties on aspect ratio with a "fast" one.
_FULL_TIER = {label for label in CANVASES if "full" in label}


def pick_canvas(image_path: str) -> str:
    """The canvas whose aspect ratio is closest to the uploaded image's own, so the framing isn't cropped/zoomed."""
    from PIL import Image

    width, height = Image.open(image_path).size
    image_ratio = width / height

    def score(label):
        canvas_height, canvas_width = CANVASES[label]
        ratio_diff = abs((canvas_width / canvas_height) - image_ratio)
        return (ratio_diff, 0 if label in _FULL_TIER else 1)

    return min(CANVASES, key=score)
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
CANVAS_MULTIPLE = 32
STEPS = 28
DEFAULT_PROMPT = (
    "The character speaks to camera in a quiet room, lips matching every word. "
    "Static camera, no zoom, no pan, no dolly movement. Keep the exact same framing, "
    "composition and distance from the subject as the reference image throughout the entire video."
)
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
REFERENCE_IMAGE_SHORT_EDGE = 2048


def snap_frames(seconds: float) -> int:
    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 = 2.0) -> None:
    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:
    return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2


def target_rows(height: int, width: int, num_frames: int) -> int:
    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(image_path: str, audio_seconds: float | None, num_frames: int) -> int:
    from PIL import Image

    width, height = Image.open(image_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)
    if audio_seconds is not None:
        rows += round(min(audio_seconds, num_frames / FPS) * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
    return rows


def get_duration(prompt_embeds, text_token_tags, image_path, audio_seconds, audio_path, height, width, num_frames, seed, **_):
    sequence = int(text_token_tags.shape[0]) + reference_rows(image_path, audio_seconds, num_frames) + target_rows(
        height, width, num_frames
    )
    denoise = STEPS * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
    encode = 5 + reference_rows(image_path, audio_seconds, 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
    return max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))


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


def load_models() -> str | None:
    """Load the denoising half (transformer_ref + both autoencoders) at startup, off the GPU."""
    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 diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
        from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
        from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3Ref2VAReferenceEncoderStep
        from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
            MiniMaxH3DecodeStep,
            MiniMaxH3Ref2VACoreDenoiseStep,
            _generation_outputs,
        )
        from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks

        class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
            """Denoising half of split ref2va: no text_encoder step, prompt_embeds/text_token_tags come in as inputs."""

            model_name = "minimax-h3"
            block_classes = [
                MiniMaxH3Ref2VASetupStep,
                MiniMaxH3Ref2VAReferenceEncoderStep,
                MiniMaxH3Ref2VACoreDenoiseStep,
                MiniMaxH3AfterDenoiseStep,
                MiniMaxH3DecodeStep,
            ]
            block_names = ["setup", "reference_encoder", "denoise", "after_denoise", "decode"]

            @property
            def outputs(self):
                return _generation_outputs()

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3Ref2VAGeneratorBlocks()
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
        pipe.load_components(dtype=torch.bfloat16)

        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer_ref.set_attention_backend(ATTENTION)

        PIPE, MANAGER = pipe, manager
        print(f"[avatar] 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: `{type(error).__name__}: {error}`"
    return LOAD_ERROR


def probe(path: str) -> tuple[float | None, float | None]:
    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 build_references(image_path: str, audio_path: str):
    from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3AudioReference, MiniMaxH3ImageReference

    return [MiniMaxH3ImageReference.from_file(image_path), MiniMaxH3AudioReference.from_file(audio_path)]


def encode_remote(prompt, image_path, audio_path, canvas, num_frames):
    from gradio_client import Client, handle_file
    from safetensors import safe_open

    client = Client(CONDITIONER_SPACE)
    path, plan = client.predict(
        prompt=prompt,
        media=[handle_file(image_path), handle_file(audio_path)],
        kinds="image,audio",
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=False,
        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()


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, image_path, audio_seconds, audio_path, height, width, num_frames, seed):
    import torch

    PIPE.to("cuda")
    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        references=build_references(image_path, audio_path),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=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(image_path, audio_path, progress=gr.Progress(track_tqdm=True)):
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The model is still loading, please try again in a moment.")
    if not image_path:
        raise gr.Error("Upload a portrait image.")
    if not audio_path:
        raise gr.Error("Upload an audio clip (the voice).")

    from diffusers.utils import encode_video

    _, audio_seconds = probe(audio_path)
    if audio_seconds is None:
        raise gr.Error("That file has no audio track.")

    # Duration is derived from the audio itself (0 == "leave it to the references").
    progress(0.0, desc="Reading the image and audio ...")
    canvas = pick_canvas(image_path)
    prompt_embeds, text_token_tags, metadata = encode_remote(DEFAULT_PROMPT, image_path, audio_path, canvas, 0)
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))

    seed = random.randint(0, 2**31 - 1)
    progress(0.15, desc=f"Generating {num_frames / FPS:.1f}s talking avatar ...")
    frames, audio, sampling_rate = _generate(
        prompt_embeds, text_token_tags, image_path, audio_seconds, audio_path, height, width, num_frames, seed
    )

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


load_models()

CSS = """
.main.fillable { max-width: 720px !important; }
"""

with gr.Blocks(title="Talking Avatar") as demo:
    gr.Markdown("# Talking Avatar\nUpload a portrait and a voice clip, then press Generate.")
    image = gr.Image(label="Portrait image", type="filepath", height=280)
    audio = gr.Audio(label="Voice", type="filepath")
    run = gr.Button("Generate", variant="primary")
    result = gr.Video(label="Talking avatar")

    run.click(generate, [image, audio], result, api_name="generate")

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