app.py

"""Unified GEPARD TTS + LongCat-Video-Avatar Gradio Space.

Startup order is load-bearing for ZeroGPU:
  1. Cache/environment variables are set before any third-party imports.
  2. create_env.setup_dependencies() re-pins transformers before ML imports.
  3. spaces is imported before torch.
  4. Both model stacks are loaded at module level and moved to "cuda" once.
"""

import os

os.environ.setdefault("HF_HOME", "/tmp/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
os.environ.setdefault("GRADIO_SSR_MODE", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("OMP_NUM_THREADS", "4")

from create_env import setup_dependencies

setup_dependencies()

try:
    import spaces  # noqa: E402
except ImportError:  # Local syntax/import checks without the Space runtime.
    class _SpacesFallback:
        @staticmethod
        def GPU(*args, **kwargs):
            if args and callable(args[0]):
                return args[0]

            def _wrap(fn):
                return fn

            return _wrap

    spaces = _SpacesFallback()

import hashlib  # noqa: E402
import json  # noqa: E402
import math  # noqa: E402
import shutil  # noqa: E402
import subprocess  # noqa: E402
import sys  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
import uuid  # noqa: E402
from collections import OrderedDict  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import imageio  # noqa: E402
import numpy as np  # noqa: E402
import soundfile as sf  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402
from PIL import Image, UnidentifiedImageError  # noqa: E402

from gepard_inference.engine import AppConfig, GenerationParams, GepardEngine  # noqa: E402
from interface import MODE_CLONE, MODE_PRESET, build_theme  # noqa: E402

ROOT = Path(__file__).parent.resolve()
CONFIG_PATH = ROOT / "config.yaml"
HF_TOKEN = os.environ.get("HF_TOKEN")
WEIGHTS_DIR = Path(os.environ.get("WEIGHTS_DIR", "weights"))
BASE_DIR = WEIGHTS_DIR / "LongCat-Video"
AVATAR_DIR = WEIGHTS_DIR / "LongCat-Video-Avatar-1.5"

SAVE_FPS = 25
NUM_FRAMES = 125
VIDEO_SECONDS = NUM_FRAMES / SAVE_FPS
AUDIO_STRIDE = 1
CP_SPLIT_HW = [1, 1]
AUDIO_GUIDANCE_SCALE = 2.0

NEGATIVE_PROMPT = (
    "Close-up, Bright tones, overexposed, static, blurred details, subtitles, "
    "style, works, paintings, images, static, overall gray, worst quality, low "
    "quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly "
    "drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, "
    "fused fingers, still picture, messy background, three legs, many people in "
    "the background, walking backwards"
)

ACCEL_MODE_EXACT = "Exact 8-step"
ACCEL_MODE_DBCACHE = "DBCache fast"
ACCEL_MODE_DBCACHE_FASTER = "DBCache faster"
EXAMPLE_CACHE_VERSION = "v4"
IMAGE_ERROR_MESSAGE = "Could not open the uploaded image. Please upload a valid image file (PNG, JPG, WEBP, etc.)"

_AUDIO_EMB_CACHE = OrderedDict()
_CACHE_LIMIT = 8
_DISK_CACHE_DIR = Path(tempfile.gettempdir()) / "avatar_demo_cache"
_AUDIO_CACHE_DIR = _DISK_CACHE_DIR / "audio_emb"
_EXAMPLE_CACHE_DIR = _DISK_CACHE_DIR / "examples"
_AUDIO_CACHE_DIR.mkdir(parents=True, exist_ok=True)
_EXAMPLE_CACHE_DIR.mkdir(parents=True, exist_ok=True)


def _install_sdpa_shim() -> None:
    """Patch xformers-style calls to PyTorch SDPA.

    A local pure-PyTorch xformers shim is bundled with this Space, so this
    function works whether or not the external xformers package is installed.
    """
    import xformers.ops

    class _BDShim:
        def __init__(self, q_seqlen, kv_seqlen):
            self.q_seqlen = list(q_seqlen)
            self.kv_seqlen = list(kv_seqlen)

        @classmethod
        def from_seqlens(cls, q_seqlen, kv_seqlen):
            return cls(q_seqlen, kv_seqlen)

    xformers.ops.fmha.attn_bias.BlockDiagonalMask = _BDShim

    def _meff(q, k, v, attn_bias=None, op=None, **_):
        if attn_bias is None:
            q_ = q.transpose(1, 2).contiguous()
            k_ = k.transpose(1, 2).contiguous()
            v_ = v.transpose(1, 2).contiguous()
            return F.scaled_dot_product_attention(q_, k_, v_).transpose(1, 2)

        if isinstance(attn_bias, _BDShim):
            outs, q_off, k_off = [], 0, 0
            for q_len, k_len in zip(attn_bias.q_seqlen, attn_bias.kv_seqlen):
                q_b = q[:, q_off : q_off + q_len].transpose(1, 2).contiguous()
                k_b = k[:, k_off : k_off + k_len].transpose(1, 2).contiguous()
                v_b = v[:, k_off : k_off + k_len].transpose(1, 2).contiguous()
                outs.append(F.scaled_dot_product_attention(q_b, k_b, v_b).transpose(1, 2))
                q_off += q_len
                k_off += k_len
            return torch.cat(outs, dim=1)

        raise NotImplementedError(f"Unsupported attn_bias in SDPA shim: {type(attn_bias)}")

    xformers.ops.memory_efficient_attention = _meff
    print("[boot] installed xformers->SDPA shim", flush=True)


def _runtime_device() -> str:
    if os.environ.get("SPACES_ZERO_GPU"):
        return "cuda"
    return "cuda" if torch.cuda.is_available() else "cpu"


def _download_longcat_weights() -> None:
    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
    print(f"[boot] WEIGHTS_DIR={WEIGHTS_DIR.resolve()}", flush=True)

    if not (BASE_DIR / "vae" / "config.json").exists():
        print("[boot] downloading LongCat-Video VAE/text encoder/tokenizer", flush=True)
        snapshot_download(
            "meituan-longcat/LongCat-Video",
            local_dir=str(BASE_DIR),
            token=HF_TOKEN,
            allow_patterns=[
                "tokenizer/*",
                "text_encoder/*.safetensors",
                "text_encoder/*.json",
                "vae/*.safetensors",
                "vae/*.json",
            ],
            ignore_patterns=[
                "text_encoder/*.fp32*",
                "text_encoder/*.bin",
                "text_encoder/flax_model*",
                "text_encoder/tf_model*",
                "vae/flax_model*",
                "vae/tf_model*",
            ],
        )

    if not (AVATAR_DIR / "base_model_int8" / "config.json").exists():
        print("[boot] downloading LongCat-Video-Avatar 1.5", flush=True)
        snapshot_download(
            "meituan-longcat/LongCat-Video-Avatar-1.5",
            local_dir=str(AVATAR_DIR),
            token=HF_TOKEN,
            allow_patterns=[
                "base_model_int8/*",
                "lora/*",
                "scheduler/*",
                "vocal_separator/*",
                "whisper-large-v3/model.safetensors",
                "whisper-large-v3/*.json",
                "whisper-large-v3/*.txt",
            ],
            ignore_patterns=[
                "whisper-large-v3/model.fp32*",
                "whisper-large-v3/flax_model*",
                "whisper-large-v3/tf_model*",
                "whisper-large-v3/pytorch_model*",
            ],
        )
    print("[boot] LongCat weights ready", flush=True)


def _patch_dit_config() -> None:
    cfg_path = AVATAR_DIR / "base_model_int8" / "config.json"
    if not cfg_path.exists():
        return

    cfg = json.loads(cfg_path.read_text())
    changed = False
    for key in ("enable_flashattn2", "enable_flashattn3", "enable_bsa"):
        if cfg.get(key):
            cfg[key] = False
            changed = True
    if not cfg.get("enable_xformers"):
        cfg["enable_xformers"] = True
        changed = True
    if changed:
        cfg_path.write_text(json.dumps(cfg, indent=2))
        print("[boot] patched DiT config -> SDPA backend", flush=True)


def _file_sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _cache_get(cache: OrderedDict, key):
    value = cache.get(key)
    if value is not None:
        cache.move_to_end(key)
    return value


def _cache_put(cache: OrderedDict, key, value) -> None:
    cache[key] = value
    cache.move_to_end(key)
    while len(cache) > _CACHE_LIMIT:
        cache.popitem(last=False)


def _cache_file(namespace: Path, key) -> Path:
    key_json = json.dumps(key, sort_keys=True, separators=(",", ":"))
    return namespace / f"{hashlib.sha256(key_json.encode('utf-8')).hexdigest()}.pt"


def _load_audio_16k(path: str):
    try:
        from scipy.signal import resample_poly

        speech, sr = sf.read(path, dtype="float32", always_2d=False)
        if speech.ndim > 1:
            speech = speech.mean(axis=1)
        if sr != 16000:
            gcd = math.gcd(int(sr), 16000)
            speech = resample_poly(speech, 16000 // gcd, int(sr) // gcd).astype(np.float32)
            sr = 16000
        return np.ascontiguousarray(speech, dtype=np.float32), sr
    except Exception as exc:
        print(f"[audio] soundfile load failed, falling back to librosa: {exc}", flush=True)
        import librosa

        speech, sr = librosa.load(path, sr=16000)
        return np.ascontiguousarray(speech, dtype=np.float32), sr


def _prepare_audio_embedding(audio_path: str, progress):
    audio_hash = _file_sha256(audio_path)
    cache_key = (audio_hash, NUM_FRAMES, SAVE_FPS, AUDIO_STRIDE)
    cached = _cache_get(_AUDIO_EMB_CACHE, cache_key)
    if cached is not None:
        progress(0.26, desc="Using cached audio conditioning")
        print(f"[cache] audio embedding hit {audio_hash[:10]}", flush=True)
        return cached.to(DEVICE, non_blocking=True)

    cache_path = _cache_file(_AUDIO_CACHE_DIR, cache_key)
    if cache_path.exists():
        try:
            cached = torch.load(cache_path, map_location="cpu")
            _cache_put(_AUDIO_EMB_CACHE, cache_key, cached)
            progress(0.26, desc="Using cached audio conditioning")
            print(f"[cache] audio embedding disk hit {audio_hash[:10]}", flush=True)
            return cached.to(DEVICE, non_blocking=True)
        except Exception as exc:
            print(f"[cache] audio disk cache read failed: {exc}", flush=True)

    t0 = time.perf_counter()
    speech, sr = _load_audio_16k(audio_path)
    pad = math.ceil((NUM_FRAMES / SAVE_FPS - len(speech) / sr) * sr)
    if pad > 0:
        speech = np.concatenate([speech, np.zeros(pad, dtype=speech.dtype)])
    print(f"[timing] audio_load={time.perf_counter() - t0:.2f}s sr={sr} samples={len(speech)}", flush=True)

    progress(0.30, desc="Encoding audio")
    t0 = time.perf_counter()
    full_audio_emb = longcat_pipe.get_audio_embedding(
        speech,
        fps=SAVE_FPS * AUDIO_STRIDE,
        device=DEVICE,
        sample_rate=sr,
        model_type="avatar-v1.5",
    )
    if torch.isnan(full_audio_emb).any():
        raise gr.Error("Audio embedding contains NaN. Try shorter text or another voice.")

    indices = torch.arange(2 * 2 + 1, device=full_audio_emb.device) - 2
    center = (
        torch.arange(0, AUDIO_STRIDE * NUM_FRAMES, AUDIO_STRIDE, device=full_audio_emb.device).unsqueeze(1)
        + indices.unsqueeze(0)
    )
    center = torch.clamp(center, min=0, max=full_audio_emb.shape[0] - 1)
    audio_emb = full_audio_emb[center][None, ...].to(DEVICE)
    print(f"[timing] audio_encode={time.perf_counter() - t0:.2f}s shape={tuple(audio_emb.shape)}", flush=True)

    audio_emb_cpu = audio_emb.detach().cpu()
    _cache_put(_AUDIO_EMB_CACHE, cache_key, audio_emb_cpu)
    try:
        torch.save(audio_emb_cpu, cache_path)
    except Exception as exc:
        print(f"[cache] audio disk cache write failed: {exc}", flush=True)
    return audio_emb


def _fit_audio_to_video_duration(audio_path: str, duration: float = VIDEO_SECONDS) -> str:
    waveform, sample_rate = sf.read(audio_path, dtype="float32", always_2d=False)
    target_samples = int(round(sample_rate * duration))
    if waveform.ndim == 1:
        current_samples = waveform.shape[0]
        if current_samples < target_samples:
            waveform = np.pad(waveform, (0, target_samples - current_samples))
        else:
            waveform = waveform[:target_samples]
    else:
        current_samples = waveform.shape[0]
        if current_samples < target_samples:
            waveform = np.pad(waveform, ((0, target_samples - current_samples), (0, 0)))
        else:
            waveform = waveform[:target_samples, :]

    fitted_path = Path(tempfile.gettempdir()) / f"longcat_drive_{uuid.uuid4().hex[:10]}.wav"
    sf.write(str(fitted_path), waveform, sample_rate)
    print(
        "[audio] video driving audio "
        f"input={current_samples / sample_rate:.3f}s output={len(waveform) / sample_rate:.3f}s "
        f"sr={sample_rate} path={fitted_path}",
        flush=True,
    )
    return str(fitted_path)


def _normalise_video_frames(frames: np.ndarray) -> np.ndarray:
    frames = np.asarray(frames)
    frame_min = float(np.nanmin(frames))
    frame_max = float(np.nanmax(frames))
    print(
        f"[video] raw_frames shape={frames.shape} dtype={frames.dtype} min={frame_min:.4f} max={frame_max:.4f}",
        flush=True,
    )
    if np.issubdtype(frames.dtype, np.floating) and frame_max <= 1.5:
        frames = frames * 255.0
    frames = np.nan_to_num(frames, nan=0.0, posinf=255.0, neginf=0.0)
    frames = np.clip(frames, 0, 255).astype(np.uint8)
    print(
        f"[video] encoded_frames count={len(frames)} duration={len(frames) / SAVE_FPS:.3f}s "
        f"dtype={frames.dtype}",
        flush=True,
    )
    return frames


def _save_video_ffmpeg_fast(frames: np.ndarray, out_base: Path, audio_path: str, fps: int, quality: int = 5) -> str:
    out_base = str(out_base)
    temp_video = out_base + "-video.mp4"
    out_path = out_base + ".mp4"

    writer = imageio.get_writer(temp_video, fps=fps, codec="libx264", quality=quality)
    try:
        for frame in frames:
            writer.append_data(np.asarray(frame))
    finally:
        writer.close()

    duration = len(frames) / fps
    cmd = [
        "ffmpeg",
        "-y",
        "-loglevel",
        "error",
        "-i",
        temp_video,
        "-i",
        audio_path,
        "-t",
        f"{duration:.3f}",
        "-map",
        "0:v:0",
        "-map",
        "1:a:0",
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-b:a",
        "96k",
        "-movflags",
        "+faststart",
        out_path,
    ]
    subprocess.run(cmd, check=True)
    try:
        os.remove(temp_video)
    except OSError:
        pass
    return out_path


def _configure_dit_acceleration(acceleration: str) -> str:
    if acceleration in (ACCEL_MODE_DBCACHE, ACCEL_MODE_DBCACHE_FASTER):
        faster = acceleration == ACCEL_MODE_DBCACHE_FASTER
        longcat_pipe.dit.configure_dbcache(
            enabled=True,
            fn=1,
            bn=0,
            warmup_steps=1,
            max_cached_steps=3 if faster else 2,
            max_continuous_cached_steps=1,
            residual_diff_threshold=0.35,
            downsample_factor=4,
        )
        return "DMD2 8-step + DBCache" + (" faster" if faster else "")

    longcat_pipe.dit.configure_dbcache(enabled=False)
    return "DMD2 8-step"


def _write_wav(sample_rate: int, waveform: np.ndarray) -> str:
    audio_path = Path(tempfile.gettempdir()) / f"gepard_{uuid.uuid4().hex[:10]}.wav"
    sf.write(str(audio_path), waveform, sample_rate)
    return str(audio_path)


def _coerce_file_path(file_value) -> str | None:
    if not file_value:
        return None
    if isinstance(file_value, dict):
        return file_value.get("path") or file_value.get("name")
    if isinstance(file_value, (str, os.PathLike)):
        return str(file_value)
    name = getattr(file_value, "name", None)
    return str(name) if name else None


def _open_reference_image(image_value):
    image_path = _coerce_file_path(image_value)
    if not image_path:
        raise gr.Error("Upload a reference image.")
    try:
        return Image.open(image_path).convert("RGB"), image_path
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise gr.Error(IMAGE_ERROR_MESSAGE) from exc


def preview_reference_image(image_value):
    if not image_value:
        return None
    try:
        _, image_path = _open_reference_image(image_value)
        return image_path
    except gr.Error:
        gr.Warning(IMAGE_ERROR_MESSAGE)
        return None


def _is_valid_image_file(path: Path) -> bool:
    try:
        with Image.open(path) as image:
            image.verify()
        return True
    except Exception:
        return False


def _resolve_example_image() -> Path | None:
    local_path = ROOT / "assets" / "avatar" / "single" / "character.png"
    if local_path.exists() and _is_valid_image_file(local_path):
        return local_path
    try:
        downloaded = hf_hub_download(
            "Mike0021/avatar-demo",
            "assets/avatar/single/character.png",
            repo_type="space",
            token=HF_TOKEN,
        )
        downloaded_path = Path(downloaded)
        if _is_valid_image_file(downloaded_path):
            return downloaded_path
    except Exception as exc:
        print(f"[examples] failed to resolve example image: {exc}", flush=True)
    return None


def _estimate_duration(*args, **kwargs) -> int:
    return 240


_install_sdpa_shim()

from transformers import AutoTokenizer, UMT5EncoderModel  # noqa: E402

from longcat_video.audio_process import (  # noqa: E402
    get_audio_encoder,
    get_audio_feature_extractor,
)
from longcat_video.modules.autoencoder_kl_wan import AutoencoderKLWan  # noqa: E402
from longcat_video.modules.quantization import load_quantized_dit  # noqa: E402
from longcat_video.modules.scheduling_flow_match_euler_discrete import (  # noqa: E402
    FlowMatchEulerDiscreteScheduler,
)
from longcat_video.pipeline_longcat_video_avatar import LongCatVideoAvatarPipeline  # noqa: E402

if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

print("[boot] loading GEPARD", flush=True)
gepard_config = AppConfig.from_yaml(CONFIG_PATH)
gepard_engine = GepardEngine(gepard_config).load()
print("[boot] GEPARD ready", flush=True)

_download_longcat_weights()
_patch_dit_config()

DEVICE = _runtime_device()
TORCH_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
print(f"[boot] LongCat device={DEVICE} dtype={TORCH_DTYPE}", flush=True)

print("[boot] LongCat tokenizer + text_encoder", flush=True)
_t = time.time()
tokenizer = AutoTokenizer.from_pretrained(str(BASE_DIR), subfolder="tokenizer", torch_dtype=TORCH_DTYPE)
text_encoder = UMT5EncoderModel.from_pretrained(str(BASE_DIR), subfolder="text_encoder", torch_dtype=TORCH_DTYPE)
print(f"[boot] text_encoder loaded in {time.time() - _t:.1f}s", flush=True)

print("[boot] LongCat VAE + scheduler", flush=True)
_t = time.time()
vae = AutoencoderKLWan.from_pretrained(str(BASE_DIR), subfolder="vae", torch_dtype=TORCH_DTYPE)
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(AVATAR_DIR), subfolder="scheduler", torch_dtype=TORCH_DTYPE)
print(f"[boot] VAE+scheduler loaded in {time.time() - _t:.1f}s", flush=True)

print("[boot] LongCat INT8 DiT + DMD2 LoRA", flush=True)
_t = time.time()
dit = load_quantized_dit(str(AVATAR_DIR), subfolder="base_model_int8", cp_split_hw=CP_SPLIT_HW)
lora_path = AVATAR_DIR / "lora" / "dmd_lora.safetensors"
if lora_path.exists():
    dit.load_lora(str(lora_path), "dmd", multiplier=1.0, lora_network_dim=128, lora_network_alpha=64)
    dit.enable_loras(["dmd"])
    print("[boot] DMD2 8-step LoRA enabled", flush=True)
print(f"[boot] DiT loaded in {time.time() - _t:.1f}s", flush=True)

print("[boot] LongCat Whisper-Large-v3", flush=True)
_t = time.time()
audio_encoder = get_audio_encoder(str(AVATAR_DIR / "whisper-large-v3"), "avatar-v1.5")
audio_feature_extractor = get_audio_feature_extractor(str(AVATAR_DIR / "whisper-large-v3"), "avatar-v1.5")
print(f"[boot] Whisper loaded in {time.time() - _t:.1f}s", flush=True)

print("[boot] assembling LongCat pipeline", flush=True)
longcat_pipe = LongCatVideoAvatarPipeline(
    tokenizer=tokenizer,
    text_encoder=text_encoder,
    vae=vae,
    scheduler=scheduler,
    dit=dit,
    audio_encoder=audio_encoder,
    audio_feature_extractor=audio_feature_extractor,
    model_type="avatar-v1.5",
)
longcat_pipe.to(DEVICE)
audio_encoder.to(DEVICE, dtype=TORCH_DTYPE)
print("[boot] LongCat ready", flush=True)


@spaces.GPU(duration=1)
def _zerogpu_probe():
    return "ready"


def _generate_talking_avatar_impl(
    text: str,
    image_path: str,
    voice_mode: str,
    speaker: str,
    reference_audio: str,
    prompt: str,
    resolution: str,
    seed: int,
    temperature: float,
    max_speech_frames: int,
    repetition_penalty: float,
    repetition_window: int,
    acceleration: str,
    progress=gr.Progress(track_tqdm=True),
):
    if not (text or "").strip():
        raise gr.Error("Enter text to synthesize.")
    image, image_path = _open_reference_image(image_path)
    reference_audio = _coerce_file_path(reference_audio)

    progress(0.02, desc="Preparing voice")
    if voice_mode == MODE_PRESET:
        if not speaker:
            raise gr.Error("Choose a preset speaker.")
        ref_codes = gepard_engine.speakers.codes(speaker)
    else:
        if not reference_audio:
            raise gr.Error("Upload a reference audio clip for voice cloning.")
        try:
            ref_codes = gepard_engine.encode_reference(reference_audio)
        except Exception as exc:
            raise gr.Error(
                "Could not process the reference audio. Please upload a clear WAV, MP3, or M4A voice clip."
            ) from exc

    params = GenerationParams(
        temperature=temperature,
        top_k=gepard_config.defaults.top_k,
        cfg_scale=gepard_config.defaults.cfg_scale,
        cfg_frames=gepard_config.defaults.cfg_frames,
        stop_threshold=gepard_config.defaults.stop_threshold,
        max_frames=int(max_speech_frames),
        repetition_penalty=float(repetition_penalty),
        repetition_window=int(repetition_window),
    )

    progress(0.08, desc="Synthesizing speech")
    t_total = time.perf_counter()
    t0 = time.perf_counter()
    sample_rate, waveform = gepard_engine.synthesize(text, ref_codes, params)
    audio_path = _write_wav(sample_rate, waveform)
    print(f"[timing] tts={time.perf_counter() - t0:.2f}s sr={sample_rate} path={audio_path}", flush=True)
    video_audio_path = _fit_audio_to_video_duration(audio_path)

    audio_emb = _prepare_audio_embedding(video_audio_path, progress)

    generation_mode = _configure_dit_acceleration(acceleration)
    progress(0.40, desc=f"Generating video ({generation_mode})")
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    clean_prompt = (prompt or default_prompt).strip()

    t0 = time.perf_counter()
    with torch.inference_mode():
        output = longcat_pipe.generate_ai2v(
            image=image,
            prompt=clean_prompt,
            negative_prompt=NEGATIVE_PROMPT,
            resolution=resolution,
            num_frames=NUM_FRAMES,
            num_inference_steps=8,
            text_guidance_scale=1.0,
            audio_guidance_scale=AUDIO_GUIDANCE_SCALE,
            output_type="np",
            generator=generator,
            audio_emb=audio_emb,
            use_distill=True,
        )
    print(f"[timing] video_generate={time.perf_counter() - t0:.2f}s mode={acceleration}", flush=True)
    if acceleration in (ACCEL_MODE_DBCACHE, ACCEL_MODE_DBCACHE_FASTER):
        print(f"[dbcache] {longcat_pipe.dit.get_dbcache_stats()}", flush=True)

    progress(0.92, desc="Muxing audio and video")
    t0 = time.perf_counter()
    frames = _normalise_video_frames(output[0])
    out_base = Path(tempfile.gettempdir()) / f"avatar_{uuid.uuid4().hex[:10]}"
    video_path = _save_video_ffmpeg_fast(frames, out_base, video_audio_path, fps=SAVE_FPS, quality=5)
    print(f"[timing] mux={time.perf_counter() - t0:.2f}s total={time.perf_counter() - t_total:.2f}s", flush=True)
    print(f"[gen] audio={audio_path} video={video_path}", flush=True)
    return audio_path, video_path


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def _generate_talking_avatar_gpu(
    text: str,
    image_path: str,
    voice_mode: str,
    speaker: str,
    reference_audio: str,
    prompt: str,
    resolution: str,
    seed: int,
    temperature: float,
    max_speech_frames: int,
    repetition_penalty: float,
    repetition_window: int,
    acceleration: str,
    progress=gr.Progress(track_tqdm=True),
):
    return _generate_talking_avatar_impl(
        text,
        image_path,
        voice_mode,
        speaker,
        reference_audio,
        prompt,
        resolution,
        seed,
        temperature,
        max_speech_frames,
        repetition_penalty,
        repetition_window,
        acceleration,
        progress,
    )


def generate_talking_avatar(
    text: str,
    image_path: str,
    voice_mode: str,
    speaker: str,
    reference_audio: str,
    prompt: str,
    resolution: str,
    seed: int,
    temperature: float,
    max_speech_frames: int,
    repetition_penalty: float,
    repetition_window: int,
    acceleration: str,
    progress=gr.Progress(track_tqdm=True),
):
    if not (text or "").strip():
        gr.Warning("Enter text to synthesize.")
        return None, None
    if not _coerce_file_path(image_path):
        gr.Warning("Upload a reference image.")
        return None, None
    try:
        with Image.open(_coerce_file_path(image_path)) as image:
            image.verify()
    except (UnidentifiedImageError, OSError, ValueError):
        gr.Warning(IMAGE_ERROR_MESSAGE)
        return None, None
    if voice_mode == MODE_CLONE and not _coerce_file_path(reference_audio):
        gr.Warning("Upload a reference audio clip for voice cloning.")
        return None, None

    return _generate_talking_avatar_gpu(
        text,
        image_path,
        voice_mode,
        speaker,
        reference_audio,
        prompt,
        resolution,
        seed,
        temperature,
        max_speech_frames,
        repetition_penalty,
        repetition_window,
        acceleration,
        progress,
    )


def _toggle_voice_mode(mode: str):
    return (
        gr.update(visible=mode == MODE_PRESET),
        gr.update(visible=mode == MODE_CLONE),
    )


CUSTOM_CSS = """
:root {
  --app-radius: 14px;
  --app-card-pad: 18px;
}
.gradio-container,
.fillable:not(.fill_width) {
  width: min(100%, 1280px) !important;
  max-width: 1280px !important;
  margin-left: auto !important;
  margin-right: auto !important;
}
#generate-btn {
  font-weight: 700;
  font-size: 1.02rem;
  padding: 0.7rem 1rem;
}

/* Hero header */
.app-header {
  display: flex;
  align-items: center;
  gap: 14px;
  padding: 4px 0 16px;
  border-bottom: 1px solid var(--border-color-primary);
  margin-bottom: 18px;
}
.app-header-accent {
  width: 8px;
  height: 46px;
  border-radius: 999px;
  background: linear-gradient(180deg, #f59e0b, #ea580c 55%, #dc2626);
  flex: 0 0 auto;
}
.app-header h1 {
  margin: 0;
  font-size: 1.5rem;
  font-weight: 800;
  line-height: 1.15;
  background: linear-gradient(90deg, #f59e0b, #ea580c 60%, #dc2626);
  -webkit-background-clip: text;
  background-clip: text;
  color: transparent;
}
.app-header p {
  margin: 4px 0 0;
  color: var(--body-text-color-subdued);
  font-size: 0.95rem;
  line-height: 1.4;
}

/* Panel cards for the two-column layout */
.panel {
  background: var(--background-fill-primary);
  border: 1px solid var(--border-color-primary);
  border-radius: var(--app-radius);
  padding: var(--app-card-pad);
  height: 100%;
}
.panel-title {
  display: flex;
  align-items: center;
  gap: 9px;
  font-weight: 700;
  font-size: 1.02rem;
  margin: 0 0 12px;
  color: var(--primary_600);
}
.panel-title.muted { color: var(--body-text-color-subdued); font-weight: 600; }
.panel-title .step {
  display: inline-flex;
  align-items: center;
  justify-content: center;
  width: 22px;
  height: 22px;
  border-radius: 50%;
  background: linear-gradient(135deg, #f59e0b, #ea580c);
  color: white;
  font-size: 0.78rem;
  font-weight: 700;
  flex: 0 0 auto;
}
.section-label {
  font-size: 0.78rem;
  font-weight: 700;
  text-transform: uppercase;
  letter-spacing: 0.06em;
  color: var(--primary_600);
  margin: 14px 0 4px;
}
.section-label:first-child { margin-top: 0; }

/* Two-step pipeline indicator on the result panel */
.pipeline-steps {
  display: flex;
  flex-wrap: wrap;
  gap: 10px;
  margin: 0 0 12px;
}
.pipeline-step {
  flex: 1 1 160px;
  border: 1px solid var(--border-color-primary);
  border-radius: 12px;
  padding: 10px 12px;
  background: var(--background-fill-secondary);
}
.pipeline-step .ps-title {
  display: flex;
  align-items: center;
  gap: 8px;
  font-size: 0.85rem;
  font-weight: 700;
  color: var(--primary_600);
}
.pipeline-step .ps-title .ps-num {
  display: inline-flex;
  align-items: center;
  justify-content: center;
  width: 18px;
  height: 18px;
  border-radius: 50%;
  background: linear-gradient(135deg, #f59e0b, #ea580c);
  color: #fff;
  font-size: 0.68rem;
}
.pipeline-step .ps-desc {
  margin: 4px 0 0;
  font-size: 0.8rem;
  color: var(--body-text-color-subdued);
  line-height: 1.35;
}
.status-box {
  border: 1px dashed var(--border-color-primary);
  border-radius: 10px;
  padding: 10px 14px;
  background: var(--background-fill-secondary);
  min-height: 22px;
  line-height: 1.45;
  font-size: 0.9rem;
  color: var(--body-text-color-subdued);
  margin: 0 0 14px;
}
.result-video {
  border-radius: 12px;
  overflow: hidden;
}

/* Examples */
.example-section {
  border-top: 1px solid var(--border-color-primary);
  margin-top: 24px;
  padding-top: 18px;
  overflow-x: hidden;
}
.example-section h3 { margin: 0 0 12px; font-size: 1.08rem; }
.example-card {
  border: 1px solid var(--border-color-primary);
  border-radius: var(--app-radius);
  padding: 14px;
  display: flex;
  gap: 14px;
  align-items: center;
  flex-wrap: wrap;
}
.example-card img { object-fit: cover; border-radius: 10px; }
.footer-note {
  text-align: center;
  color: var(--body-text-color-subdued);
  font-size: 0.85rem;
  margin-top: 20px;
}
#reference-image-upload,
#reference-image-preview {
  min-height: 160px !important;
}

/* Responsive: stack panels on narrow screens */
.app-row { flex-wrap: wrap; }
@media (max-width: 820px) {
  .app-row > div { flex-basis: 100% !important; min-width: 100% !important; }
  .app-header-accent { height: 40px; }
  .app-header h1 { font-size: 1.3rem; }
  .panel { padding: 14px; }
  .pipeline-step { flex: 1 1 100%; }
  #reference-image-upload,
  #reference-image-preview {
    min-height: 140px !important;
  }
}
"""


speaker_choices = [
    (gepard_config.speaker_labels.get(name, name), name)
    for name in gepard_engine.speakers.names
]

default_prompt = "A person talks directly to the camera with natural facial expressions and small head movements."
example_image = _resolve_example_image()
example_presets = list(gepard_config.examples[:6]) if example_image else []
example_choices = [
    (
        f"{idx + 1}. {gepard_config.speaker_labels.get(ex.speaker, ex.speaker)} - {ex.text[:78]}",
        str(idx),
    )
    for idx, ex in enumerate(example_presets)
]
if example_image:
    print(f"[examples] using image {example_image}", flush=True)
else:
    print("[examples] no valid example image found; examples hidden", flush=True)


def _example_cache_paths(index: int) -> tuple[Path, Path]:
    stem = f"{EXAMPLE_CACHE_VERSION}_{index}"
    return _EXAMPLE_CACHE_DIR / f"{stem}.wav", _EXAMPLE_CACHE_DIR / f"{stem}.mp4"


def _packaged_example_paths(index: int) -> tuple[Path, Path] | None:
    stem = f"{EXAMPLE_CACHE_VERSION}_{index}"
    packaged_dir = ROOT / "assets" / "avatar" / "cached_examples"
    audio_path = packaged_dir / f"{stem}.wav"
    video_path = packaged_dir / f"{stem}.mp4"
    if audio_path.exists() and video_path.exists():
        return audio_path, video_path
    return None


def _example_outputs(index: int, audio_path: Path, video_path: Path) -> list:
    ex = example_presets[index]
    image_path = str(example_image)
    return [
        ex.text,
        image_path,
        image_path,
        MODE_PRESET,
        ex.speaker,
        None,
        default_prompt,
        "480p",
        42,
        gepard_config.defaults.temperature,
        215,
        gepard_config.defaults.repetition_penalty,
        gepard_config.defaults.repetition_window,
        ACCEL_MODE_EXACT,
        str(audio_path),
        str(video_path),
        "✅ Cached example loaded. Listen to the speech, then watch the video — or tweak the inputs and press Generate.",
    ]


def _parse_example_index(example_index: str) -> int:
    if not example_presets:
        raise gr.Error("No valid examples are available.")

    try:
        index = int(example_index)
    except (TypeError, ValueError) as exc:
        raise gr.Error("Choose an example preset.") from exc
    if index < 0 or index >= len(example_presets):
        raise gr.Error("Choose an example preset.")
    return index


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def _generate_cached_example_gpu(index: int, progress=gr.Progress(track_tqdm=True)):
    ex = example_presets[index]
    print(f"[examples] cache miss index={index} speaker={ex.speaker}", flush=True)
    audio_path, video_path = _generate_talking_avatar_impl(
        ex.text,
        str(example_image),
        MODE_PRESET,
        ex.speaker,
        None,
        default_prompt,
        "480p",
        42,
        gepard_config.defaults.temperature,
        215,
        gepard_config.defaults.repetition_penalty,
        gepard_config.defaults.repetition_window,
        ACCEL_MODE_EXACT,
        progress,
    )
    audio_cache, video_cache = _example_cache_paths(index)
    shutil.copyfile(audio_path, audio_cache)
    shutil.copyfile(video_path, video_cache)
    return _example_outputs(index, audio_cache, video_cache)


def run_cached_example(example_index: str, progress=gr.Progress(track_tqdm=True)):
    index = _parse_example_index(example_index)
    audio_cache, video_cache = _example_cache_paths(index)
    if audio_cache.exists() and video_cache.exists():
        progress(1.0, desc="Using cached example")
        print(f"[examples] cache hit index={index}", flush=True)
        return _example_outputs(index, audio_cache, video_cache)

    packaged_paths = _packaged_example_paths(index)
    if packaged_paths is not None:
        progress(1.0, desc="Using packaged cached example")
        print(f"[examples] packaged cache hit index={index}", flush=True)
        return _example_outputs(index, *packaged_paths)

    return _generate_cached_example_gpu(index, progress)

def run_generate(
    text: str,
    image_path: str,
    voice_mode: str,
    speaker: str,
    reference_audio: str,
    prompt: str,
    resolution: str,
    seed: int,
    temperature: float,
    max_speech_frames: int,
    repetition_penalty: float,
    repetition_window: int,
    acceleration: str,
    progress=gr.Progress(track_tqdm=True),
):
    """UI-level wrapper that turns the (audio, video) backend result into a
    (status, audio, video) triple for the redesigned output panel.

    No inference happens here — the ZeroGPU-decorated ``generate_talking_avatar``
    still owns all GPU work. This only formats the result for presentation.
    """
    audio, video = generate_talking_avatar(
        text,
        image_path,
        voice_mode,
        speaker,
        reference_audio,
        prompt,
        resolution,
        seed,
        temperature,
        max_speech_frames,
        repetition_penalty,
        repetition_window,
        acceleration,
        progress,
    )
    if audio and video:
        status = "✅ Done! Play the speech preview above, then watch your talking avatar video below."
        classes = ["status-box", "status-done"]
    else:
        status = "⚠️ Generation didn't finish. Complete the inputs on the left, then press Generate."
        classes = ["status-box", "status-idle"]
    return gr.update(value=status, elem_classes=classes), audio, video


with gr.Blocks(
    title="Avatar Demo",
    theme=build_theme(),
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(
        """
        <div class="app-header">
          <div class="app-header-accent"></div>
          <div>
            <h1>Avatar Demo</h1>
            <p>Type a line, upload a face, pick a voice — get a talking avatar video.</p>
          </div>
        </div>
        """
    )

    with gr.Row(elem_classes=["app-row"]):
        # ---------------- Left panel: inputs ----------------
        with gr.Column(scale=5, elem_classes=["panel"]):
            gr.HTML('<div class="panel-title"><span class="step">1</span>Create your avatar</div>')

            gr.HTML('<div class="section-label">What should they say?</div>')
            text_in = gr.Textbox(
                label="Text",
                placeholder="Type what the avatar should say...",
                lines=4,
                max_lines=8,
                show_label=False,
            )

            gr.HTML('<div class="section-label">Upload a face (portrait)</div>')
            with gr.Row():
                image_in = gr.File(
                    label="Reference image",
                    type="filepath",
                    file_count="single",
                    elem_id="reference-image-upload",
                    show_label=False,
                    height=160,
                )
                image_preview = gr.Image(
                    label="Image preview",
                    type="filepath",
                    interactive=False,
                    height=160,
                    elem_id="reference-image-preview",
                    show_download_button=False,
                    show_share_button=False,
                )

            gr.HTML('<div class="section-label">Choose a voice</div>')
            voice_mode = gr.Radio(
                choices=[MODE_PRESET, MODE_CLONE],
                value=MODE_PRESET,
                label="Voice source",
                show_label=False,
            )
            speaker_in = gr.Dropdown(
                choices=speaker_choices,
                value=gepard_engine.speakers.names[0] if gepard_engine.speakers.names else None,
                label="Preset voice",
                visible=True,
            )
            reference_audio_in = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label=f"Reference voice (up to {int(gepard_config.max_ref_seconds)}s used)",
                visible=False,
            )

            with gr.Accordion("Advanced settings", open=False):
                prompt_in = gr.Textbox(
                    label="Video prompt",
                    value=default_prompt,
                    placeholder="Describe the desired motion in plain language.",
                    lines=2,
                )
                with gr.Row():
                    resolution_in = gr.Radio(["480p", "720p"], value="480p", label="Resolution")
                    seed_in = gr.Number(value=42, precision=0, label="Seed")
                temperature_in = gr.Slider(
                    0.05,
                    1.0,
                    value=gepard_config.defaults.temperature,
                    step=0.05,
                    label="Speech temperature",
                )
                max_speech_frames_in = gr.Slider(
                    43,
                    gepard_config.defaults.max_frames,
                    value=215,
                    step=43,
                    label="Speech frame cap",
                )
                repetition_penalty_in = gr.Slider(
                    1.0,
                    1.5,
                    value=gepard_config.defaults.repetition_penalty,
                    step=0.01,
                    label="Repetition penalty",
                )
                repetition_window_in = gr.Slider(
                    0,
                    128,
                    value=gepard_config.defaults.repetition_window,
                    step=4,
                    label="Repetition window",
                )
                acceleration_in = gr.Radio(
                    [ACCEL_MODE_EXACT, ACCEL_MODE_DBCACHE, ACCEL_MODE_DBCACHE_FASTER],
                    value=ACCEL_MODE_EXACT,
                    label="Video acceleration",
                )

            generate_btn = gr.Button(
                "Generate Talking Avatar",
                variant="primary",
                elem_id="generate-btn",
                size="lg",
            )

        # ---------------- Right panel: outputs ----------------
        with gr.Column(scale=5, elem_classes=["panel"]):
            gr.HTML('<div class="panel-title muted"><span class="step">2</span>Result</div>')

            gr.HTML(
                """
                <div class="pipeline-steps">
                  <div class="pipeline-step">
                    <div class="ps-title"><span class="ps-num">1</span>Generate speech</div>
                    <p class="ps-desc">GEPARD synthesizes a voice from your text.</p>
                  </div>
                  <div class="pipeline-step">
                    <div class="ps-title"><span class="ps-num">2</span>Generate video</div>
                    <p class="ps-desc">LongCat animates the face to match the audio.</p>
                  </div>
                </div>
                """
            )

            status_md = gr.Markdown(
                "Press **Generate** to start. The progress bar shows each step live as it runs.",
                elem_classes=["status-box"],
            )

            video_out = gr.Video(
                label="Talking avatar video",
                autoplay=True,
                height=440,
                elem_classes=["result-video"],
            )
            audio_out = gr.Audio(
                label="Step 1 preview · Generated speech",
                type="filepath",
                interactive=False,
            )

    if example_choices:
        with gr.Column(elem_classes=["example-section"]):
            gr.Markdown("### Try a cached example")
            with gr.Row(elem_classes=["example-card"]):
                gr.Image(
                    value=str(example_image),
                    label="Example avatar",
                    type="filepath",
                    interactive=False,
                    height=120,
                    show_download_button=False,
                    show_share_button=False,
                )
                with gr.Column(scale=4):
                    example_select = gr.Dropdown(
                        choices=example_choices,
                        value=example_choices[0][1],
                        label="Preset",
                    )
                    run_example_btn = gr.Button(
                        "Load cached example",
                        variant="secondary",
                    )

            example_outputs = [
                text_in,
                image_in,
                image_preview,
                voice_mode,
                speaker_in,
                reference_audio_in,
                prompt_in,
                resolution_in,
                seed_in,
                temperature_in,
                max_speech_frames_in,
                repetition_penalty_in,
                repetition_window_in,
                acceleration_in,
                audio_out,
                video_out,
                status_md,
            ]
            run_example_btn.click(
                fn=run_cached_example,
                inputs=[example_select],
                outputs=example_outputs,
                api_name="run_cached_example",
            )

    gr.HTML(
        """
        <p class="footer-note">GEPARD TTS + LongCat-Video-Avatar · runs on ZeroGPU</p>
        <script>
        (() => {
          const setAccept = () => {
            document.querySelectorAll('#reference-image-upload input[type="file"]').forEach((input) => {
              input.accept = 'image/*';
            });
          };
          setAccept();
          new MutationObserver(setAccept).observe(document.body, { childList: true, subtree: true });
        })();
        </script>
        """
    )

    image_in.change(
        fn=preview_reference_image,
        inputs=[image_in],
        outputs=[image_preview],
    )
    voice_mode.change(
        fn=_toggle_voice_mode,
        inputs=[voice_mode],
        outputs=[speaker_in, reference_audio_in],
        show_progress="hidden",
    )
    generate_btn.click(
        fn=generate_talking_avatar,
        inputs=[
            text_in,
            image_in,
            voice_mode,
            speaker_in,
            reference_audio_in,
            prompt_in,
            resolution_in,
            seed_in,
            temperature_in,
            max_speech_frames_in,
            repetition_penalty_in,
            repetition_window_in,
            acceleration_in,
        ],
        outputs=[audio_out, video_out],
    )

if __name__ == "__main__":
    demo.queue(max_size=8).launch(show_error=True)