"""
NAVA — Audio-Visual Generation on ZeroGPU.

Single-process Space wrapper around inference_nava's PromptRewriter +
ImageCaptioner + a stripped-down NAVAEngine (no SP, no dist). All heavy
weights live on CPU between calls; only @spaces.GPU-decorated functions
move them to cuda.
"""

import os
import sys
import time
import math
import importlib
import re
from pathlib import Path

import torch
import yaml
import torchaudio
from torchvision.io import write_video

# Prefer FlashAttention on H200/H100 for any F.scaled_dot_product_attention
# call NAVA's backbone makes. SDPA picks a backend dynamically; on sm_90 cards
# memory-efficient often wins by default — explicitly ordering flash first
# gives ~5-15% on attention-heavy DiT blocks. math is kept as last-resort
# fallback for shape combos flash doesn't support.
try:
    if hasattr(torch.backends.cuda, "sdp_kernel") and hasattr(
        torch.backends.cuda.sdp_kernel, "set_priority_order"
    ):
        torch.backends.cuda.sdp_kernel.set_priority_order(
            ["flash", "efficient", "math"]
        )
        print("[Setup] SDPA backend priority: flash > efficient > math")
except Exception as _e:
    print(f"[Setup] SDPA priority ordering not supported on this torch: {_e}")

import gradio as gr
import spaces
from huggingface_hub import snapshot_download
from PIL import Image

# === NAVA package (installed via requirements.txt) ===
from nava_src.utils.common import set_seed
from nava_src.models.nava.utils.model_loading_utils import load_fusion_checkpoint


_REPO = Path(__file__).resolve().parent
CONFIG_PATH = str(_REPO / "configs" / "nava.yaml")
SYSTEM_PROMPT = (_REPO / "prompts" / "rewrite_template.txt").read_text(encoding="utf-8").rstrip()


# ============================================================
# 1. Pull NAVA weights (FP8 ckpt + Wan2.2 VAE + T5 + LTX audio VAE)
# ============================================================
print("[Setup] Downloading NAVA weights from HuggingFace…")
_NAVA_DIR = snapshot_download(
    repo_id="ernie-research/NAVA",
    allow_patterns=[
        "NAVA_fp8.safetensors",          # ~13 GB FP8 ckpt
        "Wan2.2-TI2V-5B/**",             # video VAE + T5
        "params/**",                     # LTX audio VAE
        "configs/**",                    # default to vendored, but mirror for parity
    ],
)
NAVA_CKPT = os.path.join(_NAVA_DIR, "NAVA_fp8.safetensors")
print(f"[Setup] NAVA weights at {_NAVA_DIR}")

# Resolve any relative paths inside nava.yaml against the snapshot dir.
# NAVA's pipeline reads `model_id` etc. as paths; cd-ing into snapshot makes
# them resolve cleanly.
os.chdir(_NAVA_DIR)


# Module-level placeholder for the singleton classes — populated below in §6.
ENGINE = None
REWRITER = None
CAPTIONER = None


# ============================================================
# 2. Prompt Rewriter — onload/offload around each call.
#    Mirrors gradio_demo/gradio_server.py:PromptRewriter (incl. retry).
# ============================================================
_SE_PAIR_RE = re.compile(r"<S>.*?<E>", re.DOTALL)


def _count_se_pairs(text: str) -> int:
    return len(_SE_PAIR_RE.findall(text or ""))


def _extract_rewrite_lazy():
    """Late import: pe_src/rewrite.py lives in the installed NAVA package
    under repo root, not as an importable module. Inline a copy of
    extract_rewrite() so we don't depend on it being on sys.path."""
    s = """def extract_rewrite(raw):
    s = raw.strip()
    if "</think>" in s:
        s = s.rsplit("</think>", 1)[-1].strip()
    if "<think>" in s:
        s = s.split("<think>", 1)[0].strip()
    rewrite_openers = ("画面呈现", "这是一段", "这段写实", "画面中")
    looks_like_thinking = (
        "\\n" in s
        or "首先" in s[:200]
        or "分析" in s[:200]
        or "完整输出" in s
        or "改写草稿" in s
        or "最终输出" in s
        or "最终 prompt" in s
    )
    if looks_like_thinking:
        last_pos = -1
        for opener in rewrite_openers:
            pos = s.rfind(opener)
            if pos > last_pos:
                last_pos = pos
        if last_pos > 0:
            s = s[last_pos:].strip()
    end_anchors = ("整体听感", "整体氛围")
    anchor_pos = -1
    for a in end_anchors:
        p = s.rfind(a)
        if p > anchor_pos:
            anchor_pos = p
    if anchor_pos >= 0:
        tail = s[anchor_pos:]
        terminators = [tail.find(t) for t in ("。", "！", "？")]
        terminators = [t for t in terminators if t >= 0]
        if terminators:
            s = s[: anchor_pos + min(terminators) + 1].strip()
    else:
        strict_markers = (
            "注意：", "注意:", "用户说", "用户的输入", "用户没",
            "改写草稿", "最终输出", "最终 prompt", "最终prompt",
            "为了准确", "为了符合要求", "我应该", "我需要",
        )
        sentence_breaks = "。！？\\n "
        earliest = len(s)
        for m in strict_markers:
            start = 0
            while True:
                p = s.find(m, start)
                if p < 0:
                    break
                if p == 0 or s[p - 1] in sentence_breaks:
                    if p < earliest:
                        earliest = p
                    break
                start = p + 1
        if earliest < len(s):
            head = s[:earliest]
            cut = max(head.rfind("。"), head.rfind("！"), head.rfind("？"))
            if cut > 0:
                s = head[: cut + 1].strip()
            else:
                s = head.strip()
    s = s.replace("\\r", "").replace("\\n", "")
    return s.strip()
"""
    ns = {}
    exec(s, ns)
    return ns["extract_rewrite"]


_extract_rewrite = _extract_rewrite_lazy()


class PromptRewriter:
    """Loads a Qwen3 chat model to CPU; reload()/offload() around each call."""

    def __init__(self, model_path: str = "Qwen/Qwen3-4B-Instruct-2507"):
        print(f"[Rewriter] Loading {model_path} to CPU…")
        t0 = time.time()
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        load_kwargs = dict(trust_remote_code=True, torch_dtype=torch.bfloat16)
        try:
            import flash_attn  # noqa: F401
            load_kwargs["attn_implementation"] = "flash_attention_2"
            print("[Rewriter] Using flash_attention_2")
        except ImportError:
            pass
        self.model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs)
        self.model.eval()
        self._on_gpu = False
        print(f"[Rewriter] Loaded in {time.time() - t0:.1f}s (on CPU)")

    def reload(self):
        if not self._on_gpu:
            self.model.to("cuda:0")
            self._on_gpu = True

    def offload(self):
        if self._on_gpu:
            self.model.to("cpu")
            torch.cuda.empty_cache()
            self._on_gpu = False

    def rewrite(self, user_input: str, max_retries: int = 5):
        """Returns (result, warning). Auto-retries on <S><E> pair-count
        mismatch — same logic as gradio_demo."""
        self.reload()
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_input},
        ]
        chat = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        inputs = self.tokenizer(chat, return_tensors="pt").to(self.model.device)
        input_count = _count_se_pairs(user_input)
        print(f"[Rewriter] target <S><E> pairs: {input_count}")
        last_result = ""
        last_count = -1
        for attempt in range(max_retries):
            print(f"[Rewriter] Generating attempt {attempt+1}/{max_retries}…")
            t0 = time.time()
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs, max_new_tokens=4096,
                    temperature=0.3, top_p=0.75, top_k=20,
                    do_sample=True, repetition_penalty=1.05,
                )
            new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
            raw = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
            result = _extract_rewrite(raw)
            output_count = _count_se_pairs(result)
            print(f"[Rewriter] Done in {time.time()-t0:.1f}s "
                  f"({len(new_tokens)} tokens, <S><E>={output_count})")
            last_result = result
            last_count = output_count
            if input_count == 0 or output_count == input_count:
                return result, ""
            print(f"[Rewriter] mismatch (got {output_count}, want {input_count}) — retrying")
        warning = (f"⚠️ Speech 标签数量不匹配（已自动重试 {max_retries} 次）。"
                   f"输入 {input_count} 对 <S><E>，输出 {last_count} 对。请重新点击 Rewrite。")
        print(f"[Rewriter] WARN: {warning}")
        return last_result, warning


# ============================================================
# 3. VL Image Captioner — same onload/offload pattern.
# ============================================================
class ImageCaptioner:
    SYSTEM_PROMPT = (
        "你是一个视频生成提示词助手。用一段流畅的中文描述图片中的场景：人物外貌、"
        "动作、服装、背景环境、光线与色调、整体氛围。不要使用markdown格式、不要分条列举、"
        "不要说\"这张图\"或\"这是一张图片\"，直接描述画面内容，像在描述一段正在发生的"
        "视频场景。输出一段话，不超过150字。"
    )
    USER_INSTRUCTION = "请描述这张图片的视频场景。"

    def __init__(self, model_path: str = "Qwen/Qwen3-VL-4B-Instruct"):
        print(f"[Captioner] Loading {model_path} to CPU…")
        t0 = time.time()
        from transformers import AutoProcessor
        try:
            from transformers import AutoModelForImageTextToText as _Auto
        except ImportError:
            from transformers import AutoModelForCausalLM as _Auto
        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
        self.model = _Auto.from_pretrained(
            model_path, trust_remote_code=True, torch_dtype=torch.bfloat16,
        ).eval()
        self._on_gpu = False
        print(f"[Captioner] Loaded in {time.time()-t0:.1f}s (on CPU)")

    def reload(self):
        if not self._on_gpu:
            self.model.to("cuda:0")
            self._on_gpu = True

    def offload(self):
        if self._on_gpu:
            self.model.to("cpu")
            torch.cuda.empty_cache()
            self._on_gpu = False

    @torch.no_grad()
    def caption(self, image_path: str) -> str:
        self.reload()
        pil = Image.open(image_path).convert("RGB")
        msgs = [
            {"role": "system", "content": [{"type": "text", "text": self.SYSTEM_PROMPT}]},
            {"role": "user", "content": [
                {"type": "image", "image": pil},
                {"type": "text", "text": self.USER_INSTRUCTION},
            ]},
        ]
        text = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        inputs = self.processor(text=[text], images=[pil], return_tensors="pt").to(self.model.device)
        print(f"[Captioner] IN  image: {image_path}")
        t0 = time.time()
        out = self.model.generate(
            **inputs, max_new_tokens=256,
            do_sample=True, temperature=0.3, top_p=0.9,
        )
        new_tokens = out[0][inputs["input_ids"].shape[1]:]
        result = self.processor.decode(new_tokens, skip_special_tokens=True).strip()
        print(f"[Captioner] Done in {time.time()-t0:.1f}s ({len(new_tokens)} tokens)")
        print(f"[Captioner] OUT ({len(result)} chars): {result}")
        return result


def _compose_t2av_prompt(scene_caption: str, user_prompt: str) -> str:
    cap = (scene_caption or "").strip()
    spk = (user_prompt or "").strip()
    if not cap:
        return spk
    if not spk:
        return cap
    return f"{cap} {spk}"


# ============================================================
# 4. NAVA Engine — single-process, single-GPU, FP8.
#    Stripped down from gradio_demo/nava_engine.py: no SP, no dist,
#    no group_offload (single H200 has plenty of room with FP8).
# ============================================================
def _to01(x):
    return torch.clamp((x.float() + 1.0) / 2.0, 0.0, 1.0)


def _toWav(x):
    peak = x.abs().max().clamp(min=1e-12)
    x = x * (0.95 / peak)
    return x.clamp(-1.0, 1.0)


class NAVAEngineZero:
    """Single-GPU NAVA inference. All weights live on CPU between calls;
    only generate() (called from inside @spaces.GPU) moves them to cuda."""

    def __init__(self, config_path: str, ckpt_path: str):
        with open(config_path, "r") as f:
            self.cfg = yaml.safe_load(f)
        self.modality = self.cfg.get("modality", "audio_video")

        # nava.yaml references audio_config/video_config/joint_config as
        # `nava_src/models/nava/configs/...`. Those JSONs live inside the
        # installed `nava_src` site-packages, but pipeline_nava.py opens
        # them with a plain open() against cwd. Rewrite the cfg paths to
        # absolute file locations so the load works no matter where cwd is.
        import nava_src as _nava_src_pkg
        _NAVA_SRC_DIR = os.path.dirname(_nava_src_pkg.__file__)
        for key in ("audio_config", "video_config", "joint_config"):
            val = self.cfg.get("model", {}).get(key)
            if isinstance(val, str) and val.startswith("nava_src/"):
                abs_path = os.path.join(_NAVA_SRC_DIR, val[len("nava_src/"):])
                if os.path.exists(abs_path):
                    self.cfg["model"][key] = abs_path
                    print(f"[Engine] resolved {key} → {abs_path}")
                else:
                    print(f"[Engine] WARN: {key} = {val} not found at {abs_path}")
        set_seed(self.cfg.get("seed", 42))

        # Build pipeline on CPU
        module_path, class_name = self.cfg["pipeline"].rsplit(".", 1)
        PipelineClass = getattr(importlib.import_module(module_path), class_name)
        if "video" in self.modality and "audio" in self.modality:
            self.cfg["init_from_meta"] = True

        self.pipe = PipelineClass.create(
            model_id=self.cfg["model_id"],
            use_bf16=self.cfg["use_bf16"],
            audio_latent_ch=self.cfg["audio_latent_ch"],
            video_latent_ch=self.cfg["video_latent_ch"],
            lambda_ddpm=self.cfg["lambda_ddpm"],
            cfg=self.cfg,
            device=torch.device("cpu"),
        )

        # FP8 patch + load weights (state-dict route — modality/use_mmdit branch
        # mirrors inference_nava.py)
        from safetensors.torch import load_file
        print(f"[Engine] Loading {ckpt_path}…")
        state_dict = load_file(ckpt_path, device="cpu")
        is_fp8 = any(
            isinstance(v, torch.Tensor) and v.dtype == torch.float8_e4m3fn
            for v in state_dict.values()
        )
        if is_fp8:
            from NAVA_FP8 import patch_model_to_fp8
            n_patched = patch_model_to_fp8(self.pipe.model)
            print(f"[Engine] FP8 mode: patched {n_patched} Linear modules")
        if "video" in self.modality and "audio" in self.modality and not self.cfg.get("use_mmdit_model", False):
            load_fusion_checkpoint(self.pipe.model, checkpoint_path=ckpt_path, from_meta=True)
        else:
            missing, unexpected = self.pipe.model.load_state_dict(state_dict, strict=False)
            print(f"[Engine] missing={len(missing)} unexpected={len(unexpected)}")

        self.pipe.model.eval()
        self.pipe.model.backbone.set_rope_params()

        # ── dtype fix ────────────────────────────────────────────────
        # pipeline_nava.py:42-44 casts wan_vae.model weights to bf16. The VAE's
        # encode/decode paths assume an outer torch.autocast(bf16) is active so
        # that conv3d's fp32 input gets promoted at the kernel boundary. NAVA's
        # torchrun path satisfies this; under @spaces.GPU the autocast scope is
        # disabled, so conv3d sees fp32 input vs bf16 weight and crashes with
        # `Input type (float) and bias type (c10::BFloat16)…`.
        # Fix: monkey-patch both entry points (wrapped_encode for first-frame
        # latents on I2V/timbre paths, wrapped_decode for sample output) to
        # cast their input to bf16. Single chokepoint covers everything —
        # tiled_decode is reached only via wrapped_decode so we don't need to
        # patch it separately. Bf16 here matches every NAVA inference script
        # (inference_fp8.sh, inference.sh, etc.) and is what NAVA's training /
        # released videos use; bf16 vs fp32 VAE PSNR is >40dB on this model.
        if hasattr(self.pipe, "video_vae") and hasattr(self.pipe.video_vae, "wan_vae"):
            wan_vae = self.pipe.video_vae.wan_vae
            wan_vae.dtype = torch.bfloat16  # match the bf16 weight dtype

            def _force_bf16(t):
                if isinstance(t, torch.Tensor) and t.dtype != torch.bfloat16:
                    return t.to(torch.bfloat16)
                return t

            _orig_wrapped_encode = wan_vae.wrapped_encode

            def _wrapped_encode_bf16(video, *args, **kwargs):
                return _orig_wrapped_encode(_force_bf16(video), *args, **kwargs)

            wan_vae.wrapped_encode = _wrapped_encode_bf16

            _orig_wrapped_decode = wan_vae.wrapped_decode

            def _wrapped_decode_bf16(zs, *args, **kwargs):
                return _orig_wrapped_decode(_force_bf16(zs), *args, **kwargs)

            wan_vae.wrapped_decode = _wrapped_decode_bf16

        # Pipeline-internal flags read by pipe.sample()
        self.pipe._t5_offload = True   # t5 → CPU between text encodes
        self.pipe._group_offload = False

        # Inference defaults
        self.fps = self.cfg["data"].get("video_fps", 24)
        self.audio_tokens_per_sec = self.cfg["data"].get("audio_tokens_per_sec", 25)
        self.video_latent_ch = self.cfg["video_latent_ch"]
        self.patch_size = self.cfg.get("spatial_downsample", 16)
        self.dtype = torch.bfloat16 if self.cfg["use_bf16"] else torch.float16

        # Stay on CPU until generate() pulls us to cuda
        self._on_gpu = False
        print("[Engine] Ready (on CPU; will move to cuda on each generate())")

    def reload(self):
        """Move full pipeline (DiT + VAE + T5) to cuda for one inference.
        torch.compile was tried but consistently regressed throughput on this
        Space (NAVA's forward has CPU-tensor branches that force cudagraph
        skip, and inductor warmup on every cuda context reset wasn't amortized
        over a single generate)."""
        if not self._on_gpu:
            self.pipe = self.pipe.to("cuda:0")
            self._on_gpu = True

        # Wan2_2_VAE / LocalVideoVAEAdapter / LocalAudioVAEAdapter are plain
        # Python classes, not nn.Module subclasses — pipe.to(cuda) DOES NOT
        # walk into them. Their .device attr was frozen to whatever device=
        # was passed at __init__ (we passed "cpu"). Move the wrapped nn.Module
        # weights to cuda explicitly here. Without this, the entire VAE decode
        # silently runs on CPU and looks like a hang.
        try:
            wv = self.pipe.video_vae.wan_vae
            wv.model.to("cuda:0")
            wv.device = "cuda:0"
            # scale tensors used inside model.decode also need to be on GPU.
            if isinstance(wv.scale, list):
                wv.scale = [s.to("cuda:0") if isinstance(s, torch.Tensor) else s
                            for s in wv.scale]
            elif isinstance(wv.scale, torch.Tensor):
                wv.scale = wv.scale.to("cuda:0")
            print("[Engine] wan_vae.model → cuda:0")
        except Exception as e:
            print(f"[Engine] WARN failed to move wan_vae to cuda: {e}")

        try:
            av = self.pipe.audio_vae
            # Try a few common attr names — LocalAudioVAEAdapter wraps LTX VAE
            for inner_attr in ("ltx_vae", "model", "spk_model"):
                inner = getattr(av, inner_attr, None)
                if inner is None:
                    continue
                if hasattr(inner, "to"):
                    inner.to("cuda:0")
                if hasattr(inner, "device"):
                    inner.device = "cuda:0"
            print("[Engine] audio_vae components → cuda:0")
        except Exception as e:
            print(f"[Engine] WARN failed to move audio_vae to cuda: {e}")

    def offload(self):
        """Move full pipeline back to CPU between calls. ZeroGPU also tears
        down the cuda context after the @spaces.GPU function returns, so this
        is mostly a courtesy — but explicit cleanup keeps the next call's
        first .to('cuda') deterministic."""
        if self._on_gpu:
            self.pipe = self.pipe.to("cpu")
            torch.cuda.empty_cache()
            self._on_gpu = False

    # ------- batch building -------

    def _get_first_frame(self, image_path: str, height: int, width: int):
        return self.pipe.video_vae.encode(
            image_path, target_height=height, target_width=width
        ).latent_dist.sample()

    def _get_spk_embs(self, spk_wav_paths):
        out = []
        for wav_path in spk_wav_paths:
            if not wav_path or not os.path.exists(wav_path):
                out.append(torch.zeros((1, 192), dtype=torch.float32))
                continue
            query = {"data_path": wav_path, "use_spk_emb": True}
            r = self.pipe.audio_vae.encode(query).latent_dist.sample()
            out.append(r["spk_embs"])
        return out

    def _build_batch(self, prompt, image_path, spk_wav_paths, is_i2v,
                     height, width, frames):
        h = height // self.patch_size
        w = width // self.patch_size
        video_duration = ((frames - 1) * 4 + 1) / self.fps
        audio_len = math.ceil(video_duration * self.audio_tokens_per_sec)

        video_latents = torch.randn((frames, h, w, 48))
        audio_latents = torch.randn((audio_len, 48))

        first_frames = None
        if is_i2v and image_path and os.path.exists(image_path):
            first_frames = self._get_first_frame(image_path, height, width)
            video_latents = torch.randn((frames, first_frames.shape[1], first_frames.shape[2], 48))

        spk_embs = None
        if spk_wav_paths:
            spk_embs = [self._get_spk_embs(spk_wav_paths)]

        # Idempotent <extra_id_2> normalization (matches gradio_demo/nava_engine.py)
        prompt = prompt.replace("<extra_id_2>", "").replace("<S>", "<S><extra_id_2>")

        return {
            "idx": 0,
            "video_latents": video_latents,
            "first_frames": first_frames,
            "audio_latents": audio_latents,
            "save_path": ["zero_output.mp4"],
            "captions": prompt,
            "spk_embs": spk_embs,
            "is_i2v": is_i2v,
        }

    def _collate_single(self, sample):
        """Minimal single-sample collate: wrap scalars in lists, leave tensors."""
        from nava_src.data.t2v import collate_fn
        return collate_fn([sample])

    # ------- the actual inference -------

    def generate(self, prompt, image_path, spk_wav_paths, is_i2v,
                 height, width, frames, steps,
                 video_cfg, audio_cfg,
                 video_align_cfg, audio_align_cfg,
                 align_3d_cfg, timbre_cfg, timbre_align_cfg,
                 vae_tile_size=(22, 40), vae_tile_stride=(14, 26)):
        device = torch.device("cuda:0")
        # Random seed per request — single-GPU so no need to broadcast
        seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
        print(f"[Engine] seed={seed} steps={steps} {width}x{height} frames={frames}")
        set_seed(seed)

        sample = self._build_batch(prompt, image_path, spk_wav_paths, is_i2v,
                                   height, width, frames)
        batch = self._collate_single(sample)
        batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v)
                 for k, v in batch.items()}

        amp_ctx = torch.autocast(device_type="cuda", dtype=self.dtype)
        self.reload()

        # Confirm where each VAE component actually sits — silent CPU fallback
        # was the cause of "sampling 100% then hangs" before the reload() fix.
        try:
            wv_dev = next(self.pipe.video_vae.wan_vae.model.parameters()).device
            print(f"[Engine] wan_vae.model device after reload: {wv_dev}")
        except Exception as e:
            print(f"[Engine] WARN couldn't read wan_vae device: {e}")

        print(f"[Engine] sampling start: steps={steps} dtype={self.dtype}")
        t_sample0 = time.time()
        with amp_ctx:
            gen_vid_out, gen_aud_out = self.pipe.sample(
                batch,
                num_steps=steps,
                audio_guidance_scale=audio_cfg,
                video_guidance_scale=video_cfg,
                align_3d_cfg=align_3d_cfg,
                audio_align_guidance_scale=audio_align_cfg,
                video_align_guidance_scale=video_align_cfg,
                save_vid_latent=False,
                is_i2v=is_i2v,
                timbre_cfg=timbre_cfg,
                timbre_align_guidance_scale=timbre_align_cfg,
                offload_backbone=True,
                vae_cpu_offload=False,
                # H200 80GB has plenty of headroom; non-tiled VAE skips the
                # 8x conv launches + CPU-side mask blend, ~1.5x faster decode.
                tiled_vae=False,
                vae_tile_size=tuple(vae_tile_size),
                vae_tile_stride=tuple(vae_tile_stride),
                decode=True,
            )
        # pipe.sample() runs the diffusion loop AND the VAE decode internally
        # (decode=True). If we still see a long gap between this print and the
        # next one, the bottleneck is in our post-processing/write_video, not
        # in NAVA itself.
        print(f"[Engine] pipe.sample() returned in {time.time()-t_sample0:.1f}s")
        print(f"[Engine]   gen_vid_out type={type(gen_vid_out).__name__} "
              f"shape={tuple(gen_vid_out.shape) if hasattr(gen_vid_out, 'shape') else 'n/a'}")
        print(f"[Engine]   gen_aud_out type={type(gen_aud_out).__name__} len="
              f"{len(gen_aud_out) if hasattr(gen_aud_out, '__len__') else 'n/a'}")

        # Decode → mp4
        t_post0 = time.time()
        gen_vids = _to01(gen_vid_out).float()
        video_tensor = (gen_vids[0] * 255).clamp(0, 255).to(torch.uint8)
        video_tensor = video_tensor.permute(0, 2, 3, 1)
        aud = gen_aud_out[0]
        waveform = _toWav(aud["waveform"])
        if waveform.dim() == 1:
            waveform = waveform.unsqueeze(0)
        sample_rate = aud["sample_rate"]
        print(f"[Engine] post-process tensors prepped in {time.time()-t_post0:.1f}s "
              f"(video {tuple(video_tensor.shape)}, audio {tuple(waveform.shape)} @ {sample_rate}Hz)")

        out_dir = "/tmp/nava_outputs"
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, f"output_{int(time.time()*1000)}.mp4")
        t_mp4_0 = time.time()
        write_video(
            out_path, video_tensor.cpu(),
            fps=self.fps,
            video_codec="h264",
            audio_array=waveform.cpu().float().contiguous(),
            audio_fps=sample_rate,
            audio_codec="aac",
            options={"crf": "18"},
        )
        print(f"[Engine] write_video → {out_path} in {time.time()-t_mp4_0:.1f}s")
        return out_path




# ============================================================
# 5. Construct singletons (CPU only — runs once at module import)
# ============================================================
print("[Setup] Loading rewriter, captioner, NAVA engine to CPU…")
REWRITER = PromptRewriter(model_path="Qwen/Qwen3-4B-Instruct-2507")
CAPTIONER = ImageCaptioner(model_path="Qwen/Qwen3-VL-4B-Instruct")
ENGINE = NAVAEngineZero(config_path=CONFIG_PATH, ckpt_path=NAVA_CKPT)
print("[Setup] All models loaded; UI starting.")


# ============================================================
# 6. Aspect-ratio map + autodetect (mirrors gradio_demo)
# ============================================================
ASPECT_RATIO_MAP = {
    "16:9 (1280×704)": (704, 1280),
    "9:16 (704×1280)": (1280, 704),
    "1:1 (960×960)":   (960, 960),
}


def autodetect_aspect_ratio(image_path: str):
    if not image_path or not os.path.exists(image_path):
        return gr.update()
    try:
        w, h = Image.open(image_path).size
    except Exception as e:
        print(f"[Gradio] aspect autodetect failed: {e}")
        return gr.update()
    if w > h * 1.2:
        picked = "16:9 (1280×704)"
    elif h > w * 1.2:
        picked = "9:16 (704×1280)"
    else:
        picked = "1:1 (960×960)"
    print(f"[Gradio] aspect autodetect: {w}x{h} → {picked}")
    return picked


# ============================================================
# 7. Inference functions — wrapped with @spaces.GPU.
#    Each request gets up to `duration` seconds of GPU time.
# ============================================================
@spaces.GPU(duration=60)
def rewrite_fn(user_prompt: str, image_file: str):
    """VL caption (if image) → compose → rewrite. Returns
    (rewritten_with_extra_id_2, warning, vl_caption)."""
    if not user_prompt.strip():
        return "", "", ""

    # Strip stale markers; rewriter sees clean speech tags
    cap_in = user_prompt.replace("<extra_id_2>", "")

    scene = ""
    if image_file and os.path.exists(image_file):
        scene = CAPTIONER.caption(image_file)
        CAPTIONER.offload()
        cap_in = _compose_t2av_prompt(scene, cap_in)
        print(f"[Gradio] composed ({len(cap_in)} chars): {cap_in[:200]}…")

    rewritten, warning = REWRITER.rewrite(cap_in)
    REWRITER.offload()
    rewritten = rewritten.replace("<S>", "<S><extra_id_2>")
    return rewritten, warning, scene


@spaces.GPU(duration=330)
def infer_fn(user_prompt, rewritten_prompt, image_file,
             spk_wav_1, spk_wav_2,
             steps, duration_sec, aspect_ratio,
             video_cfg, audio_cfg,
             video_align_cfg, audio_align_cfg,
             align_3d_cfg, timbre_cfg, timbre_align_cfg):
    """One full diffusion run. Up to 5 minutes of GPU time per call."""
    final_prompt = rewritten_prompt.strip() if rewritten_prompt and rewritten_prompt.strip() else user_prompt.strip()
    if not final_prompt:
        return None
    height, width = ASPECT_RATIO_MAP.get(aspect_ratio, (704, 1280))
    # 24 fps; latent frames = duration * 6 + 1 (NAVA's mapping)
    frames = int(duration_sec) * 6 + 1
    is_i2v = bool(image_file)

    spk_paths = [p for p in [spk_wav_1, spk_wav_2] if p and os.path.exists(p)]

    # Make sure rewriter/captioner are off GPU before NAVA runs
    REWRITER.offload()
    CAPTIONER.offload()

    out = ENGINE.generate(
        prompt=final_prompt,
        image_path=image_file if image_file else None,
        spk_wav_paths=spk_paths or None,
        is_i2v=is_i2v,
        height=height, width=width, frames=frames,
        steps=int(steps),
        video_cfg=float(video_cfg),
        audio_cfg=float(audio_cfg),
        video_align_cfg=float(video_align_cfg),
        audio_align_cfg=float(audio_align_cfg),
        align_3d_cfg=bool(align_3d_cfg),
        timbre_cfg=bool(timbre_cfg),
        timbre_align_cfg=float(timbre_align_cfg),
    )
    ENGINE.offload()
    return out


# ============================================================
# 8. UI
# ============================================================
DEFAULT_DURATION = 5
DEFAULT_STEPS = 25
DEFAULT_FPS = 24


with gr.Blocks(title="NAVA Audio-Video Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"# NAVA — Audio-Video Generator (ZeroGPU)\n"
        f"Single H200 · FP8 · Default {DEFAULT_DURATION}s @ {DEFAULT_FPS}fps · {DEFAULT_STEPS} steps. "
        f"~5 minutes per request when the queue is short."
    )

    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown(
                "> **Tip:** ① type a short prompt (Chinese or English). "
                "② optionally upload a first-frame image — I2V mode auto-enables, aspect ratio auto-switches. "
                "③ click **Rewrite Prompt** — Qwen3 expands your input into the long Chinese caption NAVA was trained on, "
                "and (when an image is uploaded) Qwen3-VL captions the scene and composes it into the rewrite. "
                "Wrap any spoken line in `<S>...<E>` — the rewriter preserves these verbatim."
            )

            prompt_input = gr.Textbox(
                label="Prompt (原始输入)",
                placeholder="例如：一只巨龙在城市上空喷火\n或：男人愤怒的说<S>You really want to push me?<E>",
                lines=4,
            )

            rewrite_btn = gr.Button("Rewrite Prompt", variant="secondary")

            vl_caption_box = gr.Textbox(
                label="VL Caption (上传图片时自动生成；纯文本时为空)",
                lines=3,
                interactive=False,
            )

            rewritten_prompt = gr.Textbox(
                label="Rewritten Prompt (点击 Rewrite 后填充；不点则用原始输入)",
                lines=8,
                interactive=True,
            )

            speech_warning = gr.Textbox(
                label="Speech 检查",
                interactive=False,
            )

            gr.Markdown("### Image (optional — uploads enable I2V mode)")
            image_input = gr.Image(label="First Frame Image", type="filepath")

            gr.Markdown("### Speaker Reference (optional, max 2)")
            with gr.Row():
                spk_wav_1_input = gr.Audio(label="Speaker 1 WAV", type="filepath")
                spk_wav_2_input = gr.Audio(label="Speaker 2 WAV", type="filepath")

            steps_input = gr.Slider(10, 100, value=DEFAULT_STEPS, step=5, label="Inference Steps")
            duration_input = gr.Slider(
                2, 10, value=DEFAULT_DURATION, step=1,
                label=f"Duration (seconds, {DEFAULT_FPS} fps) — values above 6s may exceed the 330s ZeroGPU budget; 10s is very slow",
            )
            aspect_ratio_input = gr.Dropdown(
                choices=list(ASPECT_RATIO_MAP.keys()),
                value="16:9 (1280×704)",
                label="Aspect Ratio (auto-set when you upload an image)",
            )

            with gr.Accordion("Advanced CFG", open=False):
                video_cfg_input = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Video CFG")
                audio_cfg_input = gr.Slider(1.0, 10.0, value=2.0, step=0.5, label="Audio CFG")
                video_align_cfg_input = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Video Align CFG")
                audio_align_cfg_input = gr.Slider(1.0, 10.0, value=2.0, step=0.5, label="Audio Align CFG")
                align_3d_cfg_input = gr.Checkbox(value=True, label="3D Align CFG")
                timbre_cfg_input = gr.Checkbox(value=False, label="Timbre CFG (use speaker WAV identity)")
                timbre_align_cfg_input = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Timbre Align CFG")

            submit_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=2):
            video_output = gr.Video(label="Generated Video", height=480)

    # Wire callbacks
    image_input.change(
        fn=autodetect_aspect_ratio,
        inputs=[image_input],
        outputs=[aspect_ratio_input],
    )
    rewrite_btn.click(
        fn=rewrite_fn,
        inputs=[prompt_input, image_input],
        outputs=[rewritten_prompt, speech_warning, vl_caption_box],
    )
    submit_btn.click(
        fn=infer_fn,
        inputs=[prompt_input, rewritten_prompt, image_input,
                spk_wav_1_input, spk_wav_2_input,
                steps_input, duration_input, aspect_ratio_input,
                video_cfg_input, audio_cfg_input,
                video_align_cfg_input, audio_align_cfg_input,
                align_3d_cfg_input, timbre_cfg_input, timbre_align_cfg_input],
        outputs=[video_output],
    )


# gradio 6.x split show_api into queue(api_open=...). api_open=False keeps
# /api_info from being exposed (cleaner Space logs); the UI's API view link is
# governed separately and not relevant here.
demo.queue(max_size=20, api_open=False).launch()