"""MiniMax-H3 image-to-video (I2VA / FLF2VA) — slim ZeroGPU Space for protocol-generator.

Stripped from RioShiina/MiniMax-H3: one Gradio endpoint (`generate_video`) that runs
ComfyUI's `MiniMaxH3ImageToVideo` graph on the NVFP4-pruned FL2VA checkpoint. A start
frame is required; an optional last frame turns the job into FLF2VA. No T2VA, REF2VA,
keyframe guides, ControlNet, LoRA chains or MCP server — see README.md.
"""

import spaces  # noqa: F401  — must be imported before torch (ZeroGPU patches CUDA)

import asyncio
import glob
import json
import math
import os
import pathlib
import random
import shutil
import subprocess
import sys
import time
import uuid
from typing import Any

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download
from PIL import Image, ImageOps

# ── Paths / pins ─────────────────────────────────────────────────────────────

APP_DIR = pathlib.Path(__file__).resolve().parent
COMFY = APP_DIR / "ComfyUI"
COMFY_MODELS = COMFY / "models"
COMFY_INPUT = COMFY / "input"
COMFY_OUTPUT = COMFY / "output"

# ComfyUI master on 2026-09-12; requirements.txt mirrors this commit's own pins.
COMFYUI_COMMIT = os.environ.get("COMFYUI_COMMIT", "c75d8c966c29cb0392259af791f43373315b72db")

# Diffusion model conversion. The community NVFP4 file the Space started on lost prompt
# adherence against Wavespeed's open-weights H3 (same prompt, image and settings,
# 2026-10-16); Comfy-Org's own pruned int8_convrot is the file its I2V template
# recommends. "pruned" drops the precomputable AdaLN branches, not quality. The model
# card says to prefer int8_convrot where torch has cu130 and to use fp8_scaled only
# when it can't — switch with the Space variable H3_DIFFUSION_MODEL (int8 | fp8 | nvfp4).
DIFFUSION_MODELS = {
    "int8": ("minimax_h3_fl2va_pruned_int8_convrot.safetensors",
             "Comfy-Org/MiniMax-H3", "diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors"),
    "fp8": ("minimax_h3_fl2va_pruned_fp8_scaled.safetensors",
            "Comfy-Org/MiniMax-H3", "diffusion_models/minimax_h3_fl2va_pruned_fp8_scaled.safetensors"),
    "nvfp4": ("minimax_h3_fl2va_pruned_nvfp4.safetensors",
              "lilcheaty/MiniMax-H3-NVFP4", "minimax_h3_fl2va_pruned_nvfp4.safetensors"),
}
DIFFUSION_MODEL = os.environ.get("H3_DIFFUSION_MODEL", "int8").strip().lower()
if DIFFUSION_MODEL not in DIFFUSION_MODELS:
    print(f"[setup] unknown H3_DIFFUSION_MODEL={DIFFUSION_MODEL!r}, using int8", flush=True)
    DIFFUSION_MODEL = "int8"

# (models/<subdir>, local filename, HF repo, path in repo)
DOWNLOADS = [
    ("diffusion_models", *DIFFUSION_MODELS[DIFFUSION_MODEL]),
    ("text_encoders", "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
     "Comfy-Org/MiniMax-H3", "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"),
    ("vae", "minimax_h3_video_vae_fp16.safetensors",
     "Comfy-Org/MiniMax-H3", "vae/minimax_h3_video_vae_fp16.safetensors"),
    ("vae", "minimax_h3_audio_vae_fp32.safetensors",
     "Comfy-Org/MiniMax-H3", "vae/minimax_h3_audio_vae_fp32.safetensors"),
]
UNET_NAME = DOWNLOADS[0][1]
CLIP_NAME = DOWNLOADS[1][1]
VIDEO_VAE_NAME = DOWNLOADS[2][1]
AUDIO_VAE_NAME = DOWNLOADS[3][1]

# Override with Space secret/variable H3_JOB_DATASET. Requires HF_TOKEN (write) on the Space.
H3_JOB_DATASET = os.environ.get("H3_JOB_DATASET", "igitagoz/h3-job-outputs").strip()



def _flag(name: str, default: str) -> bool:
    return os.environ.get(name, default).strip().lower() in ("1", "true", "yes")


# Load the four checkpoints into CPU RAM once at startup so the forked GPU worker
# inherits them instead of re-reading ~34 GB from disk inside every job's GPU window.
# On 2026-10-16 that read, plus the text encode, took 94 s of a 248 s job — 38% of the
# quota. Default on; H3_PRELOAD_MODELS=0 goes back to loading inside the window.
PRELOAD_MODELS = _flag("H3_PRELOAD_MODELS", "1")

# Go one step further and place the preloaded weights on `cuda` at module level, the
# way ZeroGPU's docs say to (the CUDA emulation outside @spaces.GPU holds them and
# restores them on allocation — this is what the Wan Space does with its diffusers
# pipeline). Here it means asking ComfyUI's model management to load the four patchers
# in the parent process. Opt-in until it has been seen to work with ComfyUI under the
# emulation; a failure falls back to the CPU preload.
GPU_RESIDENT = PRELOAD_MODELS and _flag("H3_GPU_RESIDENT", "0")

# SageAttention is quantised attention (INT8 QK) — faster, near-lossless, not bit-exact.
# Off unless asked for, so the default output stays ComfyUI's reference attention.
SAGE_ATTENTION = _flag("H3_SAGE_ATTENTION", "0")

# ── Model geometry (mirrored in src/shared/pipelineCost.ts — keep in lockstep) ──

FPS = 24
CANVAS_MULTIPLE = 32
MIN_SECONDS = 1.0
MAX_SECONDS = 15.0
MIN_STEPS = 1
MAX_STEPS = 50
MIN_MAX_DIM = 256
MAX_MAX_DIM = 1920

# ZeroGPU declaration = kill timeout. Calibrated on 2026-10-16 from a 960×640 / 124
# frames / 20 steps job that held the GPU for 248 s: 94 s before the first sampling
# step (worker start, checkpoint load, text encode) and 151 s for the 20 steps plus
# the video + audio VAE decode. The old anchor (upstream's "544p / 20 steps / 5 s fits
# in 120 s") declared 205 s for that job. BASE covers the load phase while the models
# are loaded inside the window; once the preload is confirmed to cut it, lower BASE —
# and `estimateH3SpaceAllocSecsFromDims` in src/shared/pipelineCost.ts in lockstep.
# The app fits a residual on top for the *cost* figure.
MIN_GPU_SECONDS = int(os.environ.get("MIN_GPU_SECONDS", "60"))
MAX_GPU_SECONDS = int(os.environ.get("MAX_GPU_SECONDS", "1800"))
H3_REF_WORK = 960 * 544 * 124
H3_BASE_SECS = 100.0
H3_PER_STEP_SECS = 6.5
H3_WORK_EXPONENT = 1.3
H3_SAFETY = 1.25


def frame_count(seconds: float) -> int:
    """24 fps frame count snapped *up* to the model's 17k+5 grid: 5, 22, …, 124 (≈5s), 362 (≈15s)."""
    raw = int(round(float(seconds) * FPS))
    if raw <= 5:
        return 5
    return 5 + 17 * math.ceil((raw - 5) / 17)


def fit_canvas(image_w: int, image_h: int, max_dim: int) -> tuple[int, int]:
    """Longest side = max_dim, short side follows the input aspect ratio, both snapped to 32."""
    long_side = max(CANVAS_MULTIPLE, int(max_dim))
    scale = long_side / max(image_w, image_h)

    def snap(v: float) -> int:
        return max(CANVAS_MULTIPLE, round(v * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)

    return snap(image_w), snap(image_h)


def estimate_gpu_seconds(width: int, height: int, frames: int, steps: int) -> int:
    work = (width * height * frames) / H3_REF_WORK
    est = (H3_BASE_SECS + H3_PER_STEP_SECS * steps * work ** H3_WORK_EXPONENT) * H3_SAFETY
    return max(MIN_GPU_SECONDS, int(math.ceil(est)))


# ── SageAttention (optional; wheel from RioShiina's ZeroGPU build) ───────────

if SAGE_ATTENTION and "--use-sage-attention" not in sys.argv:
    sys.argv.append("--use-sage-attention")

_sage_patched = False


def apply_sage_attention_patch() -> str:
    global _sage_patched
    if not SAGE_ATTENTION:
        return "sage attention off (H3_SAGE_ATTENTION=0) — default attention"
    if _sage_patched:
        return "sage attention already enabled"
    try:
        from comfy import model_management
        import sageattention  # noqa: F401

        model_management.sage_attention_enabled = lambda: True
        model_management.pytorch_attention_enabled = lambda: False
        _sage_patched = True
        return "sage attention enabled"
    except ImportError:
        return "sageattention not installed — default attention"
    except Exception as exc:  # pragma: no cover
        return f"sage attention patch failed: {exc}"


# ── ComfyUI bootstrap ────────────────────────────────────────────────────────

def _run(cmd: list[str], cwd: pathlib.Path | None = None, check: bool = True) -> None:
    print(f"[setup] $ {' '.join(cmd)}", flush=True)
    subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=check)


def ensure_comfy() -> None:
    if not COMFY.exists():
        _run(["git", "clone", "--depth", "1", "https://github.com/comfyanonymous/ComfyUI.git", str(COMFY)])
        _run(["git", "fetch", "--depth", "1", "origin", COMFYUI_COMMIT], cwd=COMFY, check=False)
        _run(["git", "checkout", COMFYUI_COMMIT], cwd=COMFY, check=False)
    for sub in {d[0] for d in DOWNLOADS}:
        (COMFY_MODELS / sub).mkdir(parents=True, exist_ok=True)
    COMFY_INPUT.mkdir(parents=True, exist_ok=True)
    COMFY_OUTPUT.mkdir(parents=True, exist_ok=True)


def ensure_models(progress: gr.Progress | None = None) -> None:
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    for index, (sub, name, repo, repo_path) in enumerate(DOWNLOADS):
        dest = COMFY_MODELS / sub / name
        if dest.is_symlink() and not dest.exists():
            dest.unlink()
        if dest.exists():
            continue
        if progress:
            progress(index / len(DOWNLOADS), desc=f"downloading {name}")
        print(f"[setup] downloading {repo}/{repo_path}", flush=True)
        cached = hf_hub_download(repo_id=repo, filename=repo_path, token=token)
        dest.parent.mkdir(parents=True, exist_ok=True)
        os.symlink(cached, dest)


_nodes_ready = False
NODE_CLASS_MAPPINGS: dict[str, Any] = {}


def init_comfy_nodes() -> None:
    """Import ComfyUI in-process (cwd = ComfyUI/ so folder_paths resolves models/, input/, output/)."""
    global _nodes_ready, NODE_CLASS_MAPPINGS
    if _nodes_ready:
        return
    comfy_path = str(COMFY)
    sys.path = [p for p in sys.path if p != comfy_path]
    sys.path.insert(0, comfy_path)
    for module_name in list(sys.modules):
        if module_name == "utils" or module_name.startswith("utils."):
            del sys.modules[module_name]
    os.chdir(COMFY)

    import execution
    import nodes
    import server

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    server_instance = server.PromptServer(loop)
    execution.PromptQueue(server_instance)
    loop.run_until_complete(nodes.init_extra_nodes())
    NODE_CLASS_MAPPINGS = nodes.NODE_CLASS_MAPPINGS
    for required in ("MiniMaxH3ImageToVideo", "VAEDecodeAudio", "CreateVideo", "SaveVideo"):
        if required not in NODE_CLASS_MAPPINGS:
            raise RuntimeError(f"ComfyUI @ {COMFYUI_COMMIT} has no '{required}' node")
    _nodes_ready = True
    print("[setup] ComfyUI nodes ready", flush=True)


# ── Workflow (the FL2VA recipe, inlined) ────────────────────────────────────

class _DummyHidden:
    extra_pnginfo = None
    prompt = None


def node_template(class_type: str) -> dict[str, Any]:
    """Node dict pre-filled with every input's declared default, like the ComfyUI API format."""
    if class_type not in NODE_CLASS_MAPPINGS:
        raise ValueError(f"node class '{class_type}' not found")
    node_class = NODE_CLASS_MAPPINGS[class_type]
    input_types = node_class.INPUT_TYPES()
    inputs: dict[str, Any] = {}
    for name, details in {**input_types.get("required", {}), **input_types.get("optional", {})}.items():
        config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {}
        inputs[name] = config.get("default")
    return {"inputs": inputs, "class_type": class_type, "_meta": {"title": class_type}}


def build_workflow(
    *,
    prompt: str,
    width: int,
    height: int,
    frames: int,
    steps: int,
    seed: int,
    first_frame_name: str,
    last_frame_name: str | None,
    filename_prefix: str,
) -> dict[str, Any]:
    wf: dict[str, Any] = {}

    def add(node_id: str, class_type: str, title: str, **inputs: Any) -> str:
        node = node_template(class_type)
        node["_meta"]["title"] = title
        node["inputs"].update(inputs)
        wf[node_id] = node
        return node_id

    add("unet", "UNETLoader", "Load Diffusion Model", unet_name=UNET_NAME, weight_dtype="default")
    add("clip", "CLIPLoader", "Load CLIP", clip_name=CLIP_NAME, type="minimax", device="default")
    add("vvae", "VAELoader", "Load Video VAE", vae_name=VIDEO_VAE_NAME)
    add("avae", "VAELoader", "Load Audio VAE", vae_name=AUDIO_VAE_NAME)

    add("ff_load", "LoadImage", "Load First Frame", image=first_frame_name)
    # lanczos, not nearest-exact: the keyframes are usually downscaled to the canvas and
    # nearest sampling aliases them (the Wan Space resizes its keyframes with LANCZOS).
    add("ff_scale", "ImageScale", "Scale First Frame", image=["ff_load", 0],
        upscale_method="lanczos", width=width, height=height, crop="center")
    h3_inputs: dict[str, Any] = dict(
        clip=["clip", 0], vae=["vvae", 0], prompt=prompt,
        width=width, height=height, length=frames, first_frame=["ff_scale", 0],
    )
    if last_frame_name:
        add("lf_load", "LoadImage", "Load Last Frame", image=last_frame_name)
        add("lf_scale", "ImageScale", "Scale Last Frame", image=["lf_load", 0],
            upscale_method="lanczos", width=width, height=height, crop="center")
        h3_inputs["last_frame"] = ["lf_scale", 0]
    add("h3", "MiniMaxH3ImageToVideo", "MiniMax H3 Image to Video", **h3_inputs)

    add("guider", "BasicGuider", "Basic Guider", model=["unet", 0], conditioning=["h3", 0])
    add("sampler", "KSamplerSelect", "KSamplerSelect", sampler_name="res_multistep")
    add("sigmas", "BasicScheduler", "BasicScheduler", model=["unet", 0],
        scheduler="simple", steps=steps, denoise=1.0)
    add("noise", "RandomNoise", "RandomNoise", noise_seed=seed)
    add("sample", "SamplerCustomAdvanced", "SamplerCustomAdvanced", noise=["noise", 0],
        guider=["guider", 0], sampler=["sampler", 0], sigmas=["sigmas", 0], latent_image=["h3", 1])
    add("vdec", "VAEDecode", "VAE Decode", samples=["sample", 0], vae=["vvae", 0])
    add("adec", "VAEDecodeAudio", "VAE Decode Audio", samples=["sample", 0], vae=["avae", 0])
    add("video", "CreateVideo", "Create Video", images=["vdec", 0], audio=["adec", 0], fps=FPS)
    add("save", "SaveVideo", "Save Video", video=["video", 0], filename_prefix=filename_prefix,
        format="mp4", codec={"codec": "h264"})
    return wf


# ── Executor (upstream's minimal topological runner — no PromptServer round-trip) ──

def _value_at(obj: Any, index: int) -> Any:
    try:
        return obj[index]
    except (KeyError, IndexError, TypeError):
        try:
            return obj["result"][index]
        except (KeyError, IndexError, TypeError):
            return None


def _topological_order(workflow: dict[str, Any]) -> list[str]:
    from collections import defaultdict, deque

    graph: dict[str, list[str]] = defaultdict(list)
    in_degree = {node_id: 0 for node_id in workflow}
    for node_id, node in workflow.items():
        for value in node.get("inputs", {}).values():
            if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in workflow:
                graph[value[0]].append(node_id)
                in_degree[node_id] += 1
    queue = deque(n for n, d in in_degree.items() if d == 0)
    ordered: list[str] = []
    while queue:
        current = queue.popleft()
        ordered.append(current)
        for nxt in graph[current]:
            in_degree[nxt] -= 1
            if in_degree[nxt] == 0:
                queue.append(nxt)
    if len(ordered) != len(workflow):
        raise RuntimeError("workflow contains a cycle")
    return ordered


def execute_workflow(workflow: dict[str, Any], initial_objects: dict[str, Any] | None = None) -> Any:
    import torch

    computed: dict[str, Any] = dict(initial_objects or {})
    with torch.no_grad():
        order = _topological_order(workflow)
        for node_id in order:
            if node_id in computed:
                continue
            node = workflow[node_id]
            node_class = NODE_CLASS_MAPPINGS.get(node["class_type"])
            if node_class is None:
                raise RuntimeError(f"unknown node class '{node['class_type']}'")
            instance = node_class()
            for target in (node_class, instance):
                if getattr(target, "hidden", None) is None:
                    try:
                        setattr(target, "hidden", _DummyHidden())
                    except Exception:
                        pass

            kwargs: dict[str, Any] = {}
            for name, value in node["inputs"].items():
                if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in workflow:
                    src_id, out_idx = value
                    if src_id not in computed:
                        raise RuntimeError(f"output of node {src_id} needed by {node_id} but not computed")
                    actual = _value_at(computed[src_id], out_idx)
                else:
                    actual = value
                if "." in name:
                    parent, child = name.split(".", 1)
                    kwargs.setdefault(parent, {})[child] = actual
                else:
                    kwargs[name] = actual

            fn = getattr(instance, getattr(node_class, "FUNCTION"))
            print(f"[run] {node_id} ({node['class_type']})", flush=True)
            computed[node_id] = fn(**kwargs)
        return computed.get("save")


def extract_video_path(out: Any) -> str | None:
    """Locate the mp4 SaveVideo wrote (ui payload, object attrs, or newest file in output/)."""
    found: list[str] = []

    def visit(item: Any) -> None:
        if item is None or found:
            return
        if isinstance(item, (list, tuple)):
            for sub in item:
                visit(sub)
            return
        if isinstance(item, str):
            if os.path.exists(item):
                found.append(item)
            return
        if isinstance(item, dict):
            for key in ("images", "videos", "video"):
                for sub in item.get("ui", {}).get(key, []):
                    if isinstance(sub, dict) and sub.get("filename"):
                        candidate = COMFY_OUTPUT / sub.get("subfolder", "") / sub["filename"]
                        if candidate.exists():
                            found.append(str(candidate))
                            return
            return
        ui = getattr(item, "ui", None)
        if ui is not None:
            visit(ui if isinstance(ui, dict) else getattr(ui, "as_dict", lambda: None)())
        for attr in ("path", "saved_path", "filepath", "filename", "full_path"):
            val = getattr(item, attr, None)
            if isinstance(val, str):
                for candidate in (pathlib.Path(val), COMFY_OUTPUT / val):
                    if candidate.exists():
                        found.append(str(candidate))
                        return

    visit(out)
    if found:
        return found[0]
    files = glob.glob(str(COMFY_OUTPUT / "**" / "*.mp4"), recursive=True)
    if not files:
        return None
    files.sort(key=os.path.getmtime, reverse=True)
    return files[0]


# ── Preload of the four loaders (see PRELOAD_MODELS / GPU_RESIDENT) ─────────

PRELOADED: dict[str, Any] = {}
LOADER_NODES = ("unet", "clip", "vvae", "avae")


def preload_models() -> None:
    """Run the four loader nodes in the parent process (outside any GPU window)."""
    template = build_workflow(prompt="", width=960, height=544, frames=124, steps=1, seed=0,
                              first_frame_name="", last_frame_name=None, filename_prefix="")
    t0 = time.time()
    for node_id in LOADER_NODES:
        node = template[node_id]
        node_class = NODE_CLASS_MAPPINGS[node["class_type"]]
        instance = node_class()
        fn = getattr(instance, getattr(node_class, "FUNCTION"))
        print(f"[setup] preloading {node_id} ({node['class_type']})", flush=True)
        PRELOADED[node_id] = fn(**node["inputs"])
    print(f"[setup] preloaded {len(PRELOADED)} loaders into CPU RAM in {time.time() - t0:.0f}s", flush=True)


def make_models_gpu_resident() -> None:
    """Ask ComfyUI to place the preloaded patchers on the (emulated) GPU at module level.

    ZeroGPU then holds those tensors and restores them into the real GPU at allocation,
    so the worker starts with the weights in place instead of copying them up from the
    parent's RAM inside the billed window.
    """
    from comfy import model_management

    patchers = []
    for node_id in LOADER_NODES:
        obj = PRELOADED[node_id][0]
        patchers.append(getattr(obj, "patcher", obj))
    t0 = time.time()
    try:
        model_management.load_models_gpu(patchers, force_full_load=True)
    except TypeError:  # older signature without force_full_load
        model_management.load_models_gpu(patchers)
    print(f"[setup] {len(patchers)} patchers placed on cuda in {time.time() - t0:.0f}s", flush=True)


# ── Durable outputs (same scheme as the Wan / LTX Spaces) ───────────────────

def persist_job_output(job_id: str, video_path: str, seed: int) -> None:
    """Upload the finished mp4 (+ tiny sidecar) to a Hub dataset keyed by job_id so a
    client that lost the SSE stream can later download `{job_id}.mp4`."""
    job_id = (job_id or "").strip()
    if not job_id or not video_path or not os.path.isfile(video_path):
        return
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    if not token:
        print(f"[persist] skip {job_id}: no HF_TOKEN on Space", flush=True)
        return
    if not H3_JOB_DATASET:
        print(f"[persist] skip {job_id}: H3_JOB_DATASET empty", flush=True)
        return
    try:
        api = HfApi(token=token)
        api.upload_file(path_or_fileobj=video_path, path_in_repo=f"{job_id}.mp4",
                        repo_id=H3_JOB_DATASET, repo_type="dataset")
        meta = json.dumps({"job_id": job_id, "seed": int(seed)}, separators=(",", ":"))
        api.upload_file(path_or_fileobj=meta.encode("utf-8"), path_in_repo=f"{job_id}.json",
                        repo_id=H3_JOB_DATASET, repo_type="dataset")
        print(f"[persist] uploaded {job_id} → dataset {H3_JOB_DATASET}", flush=True)
    except Exception as exc:
        print(f"[persist] failed for {job_id}: {exc}", flush=True)


# ── Job helpers ──────────────────────────────────────────────────────────────

def stage_input_image(path: str, tag: str) -> tuple[str, int, int]:
    """Copy an upload into ComfyUI/input as an oriented RGB PNG; returns (name, w, h)."""
    with Image.open(path) as img:
        img = ImageOps.exif_transpose(img).convert("RGB")
        name = f"h3_{tag}_{uuid.uuid4().hex}.png"
        img.save(COMFY_INPUT / name, "PNG")
        return name, img.width, img.height


def prune_old_files(directory: pathlib.Path, max_age_secs: int) -> None:
    cutoff = time.time() - max_age_secs
    for p in directory.rglob("*"):
        try:
            if p.is_file() and p.stat().st_mtime < cutoff:
                p.unlink()
        except OSError:
            pass


def _gpu_budget(workflow: dict[str, Any], budget: int, progress: gr.Progress | None = None) -> int:
    return int(budget)


@spaces.GPU(duration=_gpu_budget)
def run_workflow_gpu(workflow: dict[str, Any], budget: int, progress: gr.Progress | None = None) -> str | None:
    # PRELOADED is read from the module global, never passed in: ZeroGPU pickles every
    # argument through a queue to the worker, and ComfyUI's model objects can't be pickled
    # ("Can't get local object 'mixed_precision_ops.<locals>.MixedPrecisionOps.Linear'",
    # 2026-10-16). The worker is forked from this process, so it inherits the global.
    print(f"[gpu] {apply_sage_attention_patch()} — budget {budget}s, "
          f"preloaded={','.join(PRELOADED) or 'none'}", flush=True)
    out = execute_workflow(workflow, PRELOADED or None)
    return extract_video_path(out)


def generate_video(
    first_frame: str | None,
    last_frame: str | None,
    prompt: str,
    max_dim: float,
    seconds: float,
    steps: float,
    seed: float,
    randomize_seed: bool,
    job_id: str = "",
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str, int, str, int]:
    """api_name=generate_video — returns (video, used_seed, job_id, duration_ms)."""
    t0 = time.time()
    if not first_frame:
        raise gr.Error("upload a start image first")
    if not str(prompt or "").strip():
        raise gr.Error("write a prompt first")

    try:
        max_dim_i = int(max_dim)
        seconds_f = float(seconds)
        steps_i = int(steps)
    except (TypeError, ValueError):
        raise gr.Error("max_dim, seconds and steps must be numbers")
    if not (MIN_MAX_DIM <= max_dim_i <= MAX_MAX_DIM):
        raise gr.Error(f"max_dim must be between {MIN_MAX_DIM} and {MAX_MAX_DIM}")
    if not (MIN_SECONDS <= seconds_f <= MAX_SECONDS):
        raise gr.Error(f"seconds must be between {MIN_SECONDS:g} and {MAX_SECONDS:g}")
    if not (MIN_STEPS <= steps_i <= MAX_STEPS):
        raise gr.Error(f"steps must be between {MIN_STEPS} and {MAX_STEPS}")

    used_seed = random.randint(0, 2**32 - 1) if randomize_seed or int(seed) < 0 else int(seed)

    progress(0.02, desc="Preparing inputs")
    prune_old_files(COMFY_INPUT, 6 * 3600)
    prune_old_files(COMFY_OUTPUT, 6 * 3600)
    ff_name, in_w, in_h = stage_input_image(first_frame, "first")
    lf_name = stage_input_image(last_frame, "last")[0] if last_frame else None

    width, height = fit_canvas(in_w, in_h, max_dim_i)
    frames = frame_count(seconds_f)
    budget = estimate_gpu_seconds(width, height, frames, steps_i)
    if budget > MAX_GPU_SECONDS:
        raise gr.Error(
            f"this job needs ~{budget}s of GPU time, above the Space cap of {MAX_GPU_SECONDS}s — "
            f"lower the resolution ({width}×{height}), duration ({frames} frames) or steps ({steps_i})"
        )
    print(f"[job] {job_id or '-'} {width}x{height} frames={frames} steps={steps_i} seed={used_seed} "
          f"last_frame={'yes' if lf_name else 'no'} budget={budget}s", flush=True)

    progress(0.05, desc="Checking models")
    ensure_models(progress)

    workflow = build_workflow(
        prompt=str(prompt), width=width, height=height, frames=frames, steps=steps_i, seed=used_seed,
        first_frame_name=ff_name, last_frame_name=lf_name,
        filename_prefix=f"video/h3_{job_id.strip() or uuid.uuid4().hex}",
    )

    progress(0.1, desc=f"Generating ({width}×{height}, {frames} frames, {steps_i} steps)")
    try:
        video = run_workflow_gpu(workflow, budget, progress=progress)
    finally:
        for name in (ff_name, lf_name):
            if name:
                try:
                    (COMFY_INPUT / name).unlink()
                except OSError:
                    pass
    if not video:
        raise gr.Error("generation finished without an output video — check the Space logs")

    persist_job_output(job_id, video, used_seed)
    duration_ms = int((time.time() - t0) * 1000)
    print(f"[job] {job_id or '-'} done in {duration_ms}ms → {video}", flush=True)
    return video, int(used_seed), job_id or "", duration_ms


# ── Startup ──────────────────────────────────────────────────────────────────

print(f"[setup] ComfyUI bootstrap — diffusion model {DIFFUSION_MODEL} ({UNET_NAME})", flush=True)
ensure_comfy()
ensure_models()
init_comfy_nodes()
print(f"[setup] {apply_sage_attention_patch()}", flush=True)
if PRELOAD_MODELS:
    try:
        preload_models()
    except Exception as exc:  # fall back to loading inside the GPU window
        PRELOADED.clear()
        print(f"[setup] preload failed, loaders will run on GPU: {exc}", flush=True)
    if PRELOADED and GPU_RESIDENT:
        try:
            make_models_gpu_resident()
        except Exception as exc:  # keep the CPU preload
            print(f"[setup] cuda placement failed, weights stay in CPU RAM: {exc}", flush=True)
else:
    print("[setup] H3_PRELOAD_MODELS=0 — loaders run inside each job's GPU window", flush=True)

# ── UI (the visible form IS the API — inputs are positional in this order) ──

with gr.Blocks(title="MiniMax-H3 I2V") as demo:
    gr.Markdown(
        "# MiniMax-H3 · image-to-video\n"
        "Start image required; add a last image for first-to-last-frame. Native audio. "
        "Slim fork of [RioShiina/MiniMax-H3](https://huggingface.co/spaces/RioShiina/MiniMax-H3) "
        "for protocol-generator — I2VA / FLF2VA only."
    )
    with gr.Row():
        first_frame = gr.Image(type="filepath", label="Start image", height=240)
        last_frame = gr.Image(type="filepath", label="Last image (optional)", height=240)
    prompt = gr.Textbox(label="Prompt", lines=4)
    with gr.Row():
        max_dim = gr.Number(label="Longest side (px, ×32)", value=960, precision=0,
                            info="960 ≈ 544p at 16:9 · 1344 ≈ 768p · 1920 ≈ 1080p")
        seconds = gr.Slider(label="Duration (s)", minimum=MIN_SECONDS, maximum=MAX_SECONDS, step=0.5, value=5.0)
        steps = gr.Slider(label="Steps", minimum=MIN_STEPS, maximum=MAX_STEPS, step=1, value=20)
    with gr.Row():
        seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
    job_id = gr.Textbox(label="job_id", value="", visible=False)
    run_btn = gr.Button("Generate", variant="primary")
    video = gr.Video(label="Result", interactive=False)
    used_seed = gr.Number(label="used_seed", precision=0, visible=False)
    job_id_out = gr.Textbox(label="job_id", visible=False)
    duration_ms = gr.Number(label="duration_ms", precision=0, visible=False)

    run_btn.click(
        fn=generate_video,
        inputs=[first_frame, last_frame, prompt, max_dim, seconds, steps, seed, randomize_seed, job_id],
        outputs=[video, used_seed, job_id_out, duration_ms],
        api_name="generate_video",
    )

demo.queue()

if __name__ == "__main__":
    demo.launch()