# spaces must be imported before torch
import spaces  # noqa: F401

import json
import hashlib
import math
import os
import shutil
import sys
import subprocess
import tempfile
import time
import uuid
import gc
from collections import OrderedDict
from pathlib import Path
from huggingface_hub import snapshot_download

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

WEIGHTS_DIR = Path(os.environ.get("WEIGHTS_DIR", "weights"))
WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
BASE_DIR = WEIGHTS_DIR / "LongCat-Video"
AVATAR_DIR = WEIGHTS_DIR / "LongCat-Video-Avatar-1.5"
print(f"[boot] WEIGHTS_DIR={WEIGHTS_DIR.resolve()}", flush=True)

sys.path.insert(0, str(Path(__file__).parent.resolve()))

import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr
from PIL import Image
import imageio

# Use persistent Hugging Face storage if available, otherwise fall back to a system temp dir
DATA_DIR = Path("/data/longcat_cache")
if not DATA_DIR.exists():
    try:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
    except Exception:
        DATA_DIR = Path(tempfile.gettempdir()) / "longcat_cache"
        DATA_DIR.mkdir(parents=True, exist_ok=True)


# ---------------------------------------------------------------------------
# Dynamic patch for an upstream bug in Meituan's official library file
# ---------------------------------------------------------------------------
def _patch_official_bug():
    target_file = Path(__file__).parent / "longcat_video" / "pipeline_longcat_video_avatar.py"
    if target_file.exists():
        try:
            content = target_file.read_text("utf-8")
            target_str = "def generate_avc("
            if target_str in content and "num_ref_latents = None" not in content:
                idx = content.find(target_str)
                body_start_idx = content.find("):", idx)
                if body_start_idx != -1:
                    # Inject a default variable into the official function body to prevent UnboundLocalError
                    patched_content = content[:body_start_idx + 2] + "\n        num_ref_latents = None" + content[body_start_idx + 2:]
                    target_file.write_text(patched_content, "utf-8")
                    print("[boot] Patched official Meituan bug (num_ref_latents = None) successfully!", flush=True)
        except Exception as e:
            print(f"[boot] Dynamic patching of official code failed: {e}", flush=True)

_patch_official_bug()


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


# ---------------------------------------------------------------------------
# 0) xformers → SDPA shim
# ---------------------------------------------------------------------------

def _install_sdpa_shim():
    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)


_install_sdpa_shim()


# ---------------------------------------------------------------------------
# 1) Download weights
# ---------------------------------------------------------------------------

token = os.environ.get("HF_TOKEN")

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=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 (INT8 + lora + whisper + vocal_separator)…", flush=True)
    snapshot_download(
        "meituan-longcat/LongCat-Video-Avatar-1.5",
        local_dir=str(AVATAR_DIR),
        token=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] weights ready", flush=True)


# ---------------------------------------------------------------------------
# 2) Patch DiT config to use SDPA backend
# ---------------------------------------------------------------------------

_cfg_path = AVATAR_DIR / "base_model_int8" / "config.json"
if _cfg_path.exists():
    _cfg = json.loads(_cfg_path.read_text())
    _changed = False
    for k in ("enable_flashattn2", "enable_flashattn3", "enable_bsa"):
        if _cfg.get(k):
            _cfg[k] = 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)


# ---------------------------------------------------------------------------
# 3) Eager model load at module level
# ---------------------------------------------------------------------------

from transformers import AutoTokenizer, UMT5EncoderModel  # noqa: E402

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

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.bfloat16 if device == "cuda" else torch.float32
CP_SPLIT_HW = [1, 1]

print(f"[boot] device={device} dtype={torch_dtype}", flush=True)
print("[boot] 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 25.1s", flush=True)

print("[boot] VAE + scheduler…", flush=True); _t = time.time()
vae = AutoencoderKLWan.from_pretrained(str(BASE_DIR), subfolder="vae", torch_dtype=torch_dtype)
vae.enable_slicing()  # Enable slicing for lighter video-memory rendering (safe, no known issues)
# vae.enable_tiling()  # Permanently disabled: tiling triggers an upstream bug on multi-frame inputs
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(AVATAR_DIR), subfolder="scheduler", torch_dtype=torch_dtype)
print(f"[boot] VAE+scheduler loaded in 0.4s", flush=True)

print("[boot] 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 96.4s", flush=True)

print("[boot] 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 4.2s", flush=True)

print("[boot] vocal separator (Kim_Vocal_2)…")
from audio_separator.separator import Separator  # noqa: E402

VOCAL_TMP = Path("/tmp/vocal_out")
(VOCAL_TMP / "vocals").mkdir(parents=True, exist_ok=True)
vocal_separator = Separator(
    output_dir=str(VOCAL_TMP / "vocals"),
    model_file_dir=str(AVATAR_DIR / "vocal_separator"),
)
vocal_separator.load_model("Kim_Vocal_2.onnx")

print("[boot] assembling pipeline…", flush=True)
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",
)
pipe.to(device)
audio_encoder.to(device, dtype=torch_dtype)
print("[boot] ready.", flush=True)


# ---------------------------------------------------------------------------
# 4) Inference Engine with Save & Resume support
# ---------------------------------------------------------------------------

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

VOCAL_MODE_FAST = "Clean speech (fast)"
VOCAL_MODE_QUALITY = "Isolate vocals (quality)"
ACCEL_MODE_EXACT = "Exact 8-step"
ACCEL_MODE_DBCACHE = "DBCache fast"
ACCEL_MODE_DBCACHE_FASTER = "DBCache faster"
SAVE_FPS = 25
_AUDIO_EMB_CACHE = OrderedDict()
_VOCAL_CACHE = OrderedDict()
_CACHE_LIMIT = 8
_DISK_CACHE_DIR = Path(tempfile.gettempdir()) / "longcat_cache"
_AUDIO_CACHE_DIR = _DISK_CACHE_DIR / "audio_emb"
_AUDIO_CACHE_DIR.mkdir(parents=True, exist_ok=True)

CUSTOM_CSS = """
main,
.gradio-container,
.fillable:not(.fill_width) {
  width: min(100%, 1320px) !important;
  max-width: 1320px !important;
  margin-left: auto !important;
  margin-right: auto !important;
}
"""


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):
    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:
        import soundfile as sf
        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 e:
        print(f"[audio] soundfile load failed, falling back to librosa: {e}", flush=True)
        import librosa

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


def _extract_vocal(src: str, audio_hash: str) -> str:
    cached_path = _cache_get(_VOCAL_CACHE, audio_hash)
    if cached_path and Path(cached_path).exists():
        print(f"[cache] vocal hit {audio_hash[:10]}", flush=True)
        return cached_path

    stable_path = VOCAL_TMP / "vocals" / f"{audio_hash[:16]}_vocals.wav"
    if stable_path.exists():
        _cache_put(_VOCAL_CACHE, audio_hash, str(stable_path))
        return str(stable_path)

    try:
        outputs = vocal_separator.separate(src)
        if outputs:
            separated = (VOCAL_TMP / "vocals" / outputs[0]).resolve()
            shutil.copyfile(separated, stable_path)
            _cache_put(_VOCAL_CACHE, audio_hash, str(stable_path))
            return str(stable_path)
    except Exception as e:
        print(f"[vocal] separation failed, using raw audio: {e}", flush=True)
    return src


def _check_duration(*args, **kwargs):
    return 120


def _prepare_audio_embedding(audio_path: str, vocal_mode: str, total_target_frames: int, save_fps: int, audio_stride: int, progress):
    audio_hash = _file_sha256(audio_path)
    key = (audio_hash, vocal_mode, total_target_frames)
    if key in _AUDIO_EMB_CACHE:
        return _AUDIO_EMB_CACHE[key]

    t0 = time.perf_counter()
    if vocal_mode == VOCAL_MODE_QUALITY:
        progress(0.05, desc="Isolating vocals…")
        vocal_path = _extract_vocal(audio_path, audio_hash)
    else:
        progress(0.05, desc="Using clean speech directly…")
        vocal_path = audio_path
    print(f"[timing] audio_input_ready={time.perf_counter() - t0:.2f}s mode={vocal_mode}", flush=True)

    t0 = time.perf_counter()
    speech, sr = _load_audio_16k(vocal_path)
    pad = math.ceil((total_target_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=16000 samples={len(speech)}", flush=True)

    progress(0.15, desc="Encoding audio (Whisper-Large-v3)…")
    t0 = time.perf_counter()
    full_audio_emb = 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 a different audio clip.")

    indices = torch.arange(2 * 2 + 1, device=full_audio_emb.device) - 2
    center = torch.arange(0, audio_stride * total_target_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_CACHE[key] = audio_emb
    return audio_emb


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",
        "-shortest",
        out_path,
    ]
    subprocess.run(cmd, check=True)
    try:
        os.remove(temp_video)
    except OSError:
        pass
    return out_path


def _encode_frames_silent_mp4(frames, out_path: str, fps: int, quality: int = 5) -> str:
    """Encode a chunk of raw frames (no audio) into a temporary, compressed silent video file."""
    writer = imageio.get_writer(out_path, fps=fps, codec="libx264", quality=quality)
    try:
        for frame in frames:
            writer.append_data(np.asarray(frame))
    finally:
        writer.close()
    return out_path


def _concat_mp4_stream_copy(video_paths, out_path: str) -> str:
    """Concatenate multiple video segments into one continuous file with no re-encoding (stream copy)."""
    video_paths = [p for p in video_paths if p and os.path.exists(p)]
    if len(video_paths) == 1:
        shutil.copyfile(video_paths[0], out_path)
        return out_path

    list_file = out_path + ".concat_list.txt"
    with open(list_file, "w", encoding="utf-8") as f:
        for p in video_paths:
            f.write(f"file '{os.path.abspath(p)}'\n")
    cmd = [
        "ffmpeg", "-y", "-loglevel", "error",
        "-f", "concat", "-safe", "0", "-i", list_file,
        "-c", "copy", out_path,
    ]
    subprocess.run(cmd, check=True)
    try:
        os.remove(list_file)
    except OSError:
        pass
    return out_path


def _mux_final_video(video_only_path: str, audio_path: str, out_path: str, total_frames: int, fps: int) -> str:
    """Mux the completed silent video track with the user's original audio file."""
    duration = total_frames / fps
    cmd = [
        "ffmpeg", "-y", "-loglevel", "error",
        "-i", video_only_path,
        "-i", audio_path,
        "-t", f"{duration:.3f}",
        "-map", "0:v:0",
        "-map", "1:a:0",
        "-c:v", "copy",
        "-c:a", "aac",
        "-b:a", "96k",
        "-shortest",
        out_path,
    ]
    subprocess.run(cmd, check=True)
    return out_path


def _chunked_vae_decode_step(pipe, denorm_latents, start_frame_idx: int, feat_map_state, conv_idx_state,
                              start_gpu_time: float, timeout: float = 95.0):
    """
    Denoise and render the final output (VAE decode) frame-by-frame, with safe pause/resume support.

    This mirrors the pipeline's internal `_decode` logic, but instead of rendering the whole video
    in one pass (which would exceed the available GPU session time on long videos and raise a
    RuntimeError), frames are rendered one at a time. If the session is about to run out of time,
    the VAE's causal cache state (feat_cache) is saved so rendering can resume in the next session
    from exactly the same point, with no skipped or discontinuous frames.
    """
    num_frame = denorm_latents.shape[2]

    if feat_map_state is not None:
        pipe.vae._feat_map = [
            (t.to(device=device, dtype=pipe.vae.dtype) if torch.is_tensor(t) else t) for t in feat_map_state
        ]
        pipe.vae._conv_idx = list(conv_idx_state)
    else:
        pipe.vae.clear_cache()

    new_frames = []
    i = start_frame_idx
    timed_out = False
    with torch.inference_mode():
        while i < num_frame:
            elapsed = time.time() - start_gpu_time
            if elapsed > timeout:
                timed_out = True
                break

            z = denorm_latents[:, :, i:i + 1].to(device=device, dtype=pipe.vae.dtype)
            x = pipe.vae.post_quant_conv(z)
            pipe.vae._conv_idx = [0]
            if i == 0:
                out = pipe.vae.decoder(x, feat_cache=pipe.vae._feat_map, feat_idx=pipe.vae._conv_idx, first_chunk=True)
            else:
                out = pipe.vae.decoder(x, feat_cache=pipe.vae._feat_map, feat_idx=pipe.vae._conv_idx)

            if pipe.vae.config.patch_size is not None:
                out = unpatchify(out, patch_size=pipe.vae.config.patch_size)
            out = torch.clamp(out, min=-1.0, max=1.0)

            frame_np = pipe.video_processor.postprocess_video(out)[0]
            frame_np = (frame_np * 255).astype(np.uint8)
            # Each latent frame can map to several pixel frames (1 for the first, 4 for later ones,
            # due to 4x temporal upsampling); each pixel frame must be appended to the list
            # individually rather than as one multi-frame array item (the latter caused a
            # ValueError in the video writer).
            for single_frame in frame_np:
                new_frames.append(single_frame)

            del out, x, z
            i += 1

    feat_map_cpu = [(t.cpu() if torch.is_tensor(t) else t) for t in pipe.vae._feat_map]
    conv_idx_cpu = list(pipe.vae._conv_idx)

    return {
        "done": i >= num_frame,
        "next_frame_idx": i,
        "feat_map": feat_map_cpu,
        "conv_idx": conv_idx_cpu,
        "frames": new_frames,
        "timed_out": timed_out,
    }


def _configure_dit_acceleration(acceleration: str):
    if acceleration in (ACCEL_MODE_DBCACHE, ACCEL_MODE_DBCACHE_FASTER):
        faster = acceleration == ACCEL_MODE_DBCACHE_FASTER
        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 "")

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


@spaces.GPU(duration=_check_duration, size="large")
def generate(
    image_path: str,
    audio_path: str,
    prompt: str,
    resolution: str,
    seed: int,
    vocal_mode: str = VOCAL_MODE_FAST,
    acceleration: str = ACCEL_MODE_DBCACHE_FASTER,
    state_file: str = None,
    progress=gr.Progress(track_tqdm=True),
):
    save_fps = SAVE_FPS
    audio_stride = 1
    start_gpu_time = time.time()
    dit_dtype = pipe.dit.dtype

    is_resume = False
    state_dict = None

    # 1. Stateless client-to-client recovery: restore and resume a previously interrupted task
    if state_file and os.path.exists(state_file):
        try:
            yield None, None, gr.Markdown("⏳ **Restoring the state file on a new GPU session...**")
            state_dict = torch.load(state_file, map_location="cpu")
            
            # Rebuild the reference image and driving audio files on the newly isolated container
            unique_prefix = uuid.uuid4().hex[:8]
            temp_img_path = os.path.join(tempfile.gettempdir(), f"{unique_prefix}_{state_dict['img_name']}")
            temp_aud_path = os.path.join(tempfile.gettempdir(), f"{unique_prefix}_{state_dict['aud_name']}")
            
            with open(temp_img_path, "wb") as f:
                f.write(state_dict["img_bytes"])
            with open(temp_aud_path, "wb") as f:
                f.write(state_dict["aud_bytes"])
                
            image_path = temp_img_path
            audio_path = temp_aud_path
            
            total_target_frames = state_dict["total_target_frames"]
            current_frame_offset = state_dict["current_frame_offset"]
            prompt = state_dict["prompt"]
            resolution = state_dict["resolution"]
            seed = state_dict["seed"]
            accumulated_latents = state_dict["accumulated_latents"]
            
            # Restore the diffusion step and any partially-finished chunk's active latents
            current_chunk_step = state_dict.get("current_chunk_step", 0)
            active_chunk_latents = state_dict.get("active_chunk_latents", None)
            
            img_bytes = state_dict["img_bytes"]
            img_name = state_dict["img_name"]
            aud_bytes = state_dict["aud_bytes"]
            aud_name = state_dict["aud_name"]

            # Restore the final-render (VAE decode) phase state, in case the previous session stopped there
            resume_phase = state_dict.get("phase", "generate")
            decode_frame_idx = state_dict.get("decode_frame_idx", 0)
            decode_feat_map = state_dict.get("decode_feat_map", None)
            decode_conv_idx = state_dict.get("decode_conv_idx", None)
            decode_video_bytes = state_dict.get("decode_video_bytes", None)
            decoded_pixel_frame_count = state_dict.get("decoded_pixel_frame_count", 0)
            if resume_phase == "decode":
                # Skip straight past the generation loop and go directly into the final render phase
                current_frame_offset = total_target_frames

            is_resume = True
            print(f"[resume] Successfully restored task from client state. Frame offset: {current_frame_offset}, Chunk step: {current_chunk_step}, Phase: {resume_phase}", flush=True)
        except Exception as e:
            print(f"[resume] Failed to load state file: {e}", flush=True)
            is_resume = False

    # If this is not a resume, start fresh and compute the frame timing
    if not is_resume:
        resume_phase = "generate"
        decode_frame_idx = 0
        decode_feat_map = None
        decode_conv_idx = None
        decode_video_bytes = None
        decoded_pixel_frame_count = 0
        if not image_path:
            raise gr.Error("Please upload a reference image.")
        if not audio_path:
            raise gr.Error("Please upload an audio clip.")
        
        # Detect the audio duration and derive the frame count using the 4n + 1 formula
        audio_duration = 5.0
        try:
            import librosa
            audio_duration = float(librosa.get_duration(path=audio_path))
        except Exception:
            audio_duration = 5.0
            
        target_duration = max(1.0, min(120.0, audio_duration))  # supports up to 120 seconds (2 minutes)
        raw_frames = int(target_duration * save_fps)
        n = round((raw_frames - 1) / 4)
        total_target_frames = 4 * n + 1
        total_target_frames = max(25, min(3001, total_target_frames))  # cap of 3000 frames (2 minutes)
        accumulated_latents = None
        current_frame_offset = 0
        current_chunk_step = 0
        active_chunk_latents = None
        
        # Read the raw image/audio bytes so they can be embedded in the client-side state file
        with open(image_path, "rb") as f:
            img_bytes = f.read()
        img_name = Path(image_path).name
        
        with open(audio_path, "rb") as f:
            aud_bytes = f.read()
        aud_name = Path(audio_path).name
        
        print(f"[init] Calculated frames: {total_target_frames} for {target_duration:.2f}s audio.", flush=True)

    # Sync and load the Whisper audio features
    audio_emb = _prepare_audio_embedding(audio_path, vocal_mode, total_target_frames, save_fps, audio_stride, progress)

    generation_mode = _configure_dit_acceleration(acceleration)
    image = Image.open(image_path).convert("RGB")
    generator = torch.Generator(device=device).manual_seed(int(seed))
    
    scale_factor_spatial = pipe.vae_scale_factor_spatial * 2
    height, width = pipe.get_condition_shape(image, resolution, scale_factor_spatial=scale_factor_spatial)

    # Chunk length is capped to keep VRAM usage in check and avoid CUDA out-of-memory errors
    max_chunk_size = 125
    # If the whole video is 165 frames (~6.6s) or shorter, render it in a single chunk
    if total_target_frames <= 165:
        first_chunk_size = total_target_frames
    else:
        first_chunk_size = 125

    # Reset the session start time precisely, to discount any Hugging Face queueing delay
    start_gpu_time = time.time()

    # 2. Autoregressive chunk-by-chunk generation loop, with session-time-budget management
    while current_frame_offset < total_target_frames:
        elapsed = time.time() - start_gpu_time
        
        # If the remaining time is running critically low, pause and package the current state
        if elapsed > 95.0:
            progress(1.0, desc="Packaging intermediate state into state file (.pt)…")
            state_to_save = {
                "accumulated_latents": accumulated_latents.cpu() if accumulated_latents is not None else None,
                "current_frame_offset": current_frame_offset,
                "total_target_frames": total_target_frames,
                "prompt": prompt,
                "resolution": resolution,
                "seed": seed,
                "img_bytes": img_bytes,
                "img_name": img_name,
                "aud_bytes": aud_bytes,
                "aud_name": aud_name,
                "current_chunk_step": current_chunk_step,
                "active_chunk_latents": active_chunk_latents.cpu() if active_chunk_latents is not None else None,
            }
            temp_state_file = os.path.join(tempfile.gettempdir(), f"avatar_state_{uuid.uuid4().hex[:6]}.pt")
            torch.save(state_to_save, temp_state_file)
            
            status_msg = (
                f"⏸️ **Processing time for this session has ended.**\n\n"
                f"📦 **Checkpoint saved:** progress up to frame **{current_frame_offset}** "
                f"(of **{total_target_frames}** total frames) has been saved to a `.pt` state file — "
                f"this is a snapshot of exactly where generation stopped, not an error.\n\n"
                f"👇 **To continue:**\n"
                f"1. **Download** the state file from the box on the right.\n"
                f"2. This file can be uploaded again at any time — in a new session here, on your own "
                f"local GPU, or anywhere else that runs this pipeline — with no loss of progress.\n"
                f"3. Upload it into the **\"Upload state file (.pt)\"** field and click Generate again: "
                f"rendering will resume from the exact frame and diffusion step it left off at."
            )
            yield None, temp_state_file, gr.Markdown(status_msg)
            return

        # Compute the next chunk's length, capped by the max chunk size
        remaining_frames = total_target_frames - current_frame_offset
        if current_frame_offset == 0:
            num_frames_chunk = first_chunk_size
        else:
            num_frames_chunk_raw = 13 + remaining_frames
            n_chunk = math.ceil((num_frames_chunk_raw - 1) / 4)
            num_frames_chunk = 4 * n_chunk + 1
            num_frames_chunk = min(max_chunk_size, num_frames_chunk)

        chunk_end = min(current_frame_offset + (num_frames_chunk - 13), total_target_frames) if current_frame_offset > 0 else num_frames_chunk
        
        status_info = (
            f"🔄 **Generating new segment:** frame **{current_frame_offset}** to **{chunk_end}** (of **{total_target_frames}** total frames) "
            f"diffusion step: {current_chunk_step}/8..."
        )
        yield None, None, gr.Markdown(status_info)

        if current_frame_offset == 0 or (current_frame_offset == first_chunk_size and current_chunk_step > 0 and accumulated_latents is None):
            # Start (or resume) phase-one generation (generate_ai2v)
            audio_emb_chunk = audio_emb[:, :first_chunk_size]
            # Important fix: latents saved mid-chunk (phase one) must be moved back to the correct
            # device/dtype before use, otherwise this raises a device/dtype mismatch RuntimeError.
            latents_input = active_chunk_latents.to(device=device, dtype=dit_dtype) if active_chunk_latents is not None else None
            
            try:
                with torch.inference_mode():
                    chunk_latents = pipe.generate_ai2v(
                        image=image,
                        prompt=prompt,
                        negative_prompt=NEGATIVE_PROMPT,
                        resolution=resolution,
                        num_frames=first_chunk_size,
                        num_inference_steps=8,
                        text_guidance_scale=1.0,
                        audio_guidance_scale=1.0,
                        output_type="latent",
                        generator=generator,
                        audio_emb=audio_emb_chunk,
                        use_distill=True,
                        latents=latents_input,
                        start_step=current_chunk_step,
                        start_time=start_gpu_time,
                        timeout=95.0,
                    )
                accumulated_latents = chunk_latents.cpu()
                current_frame_offset = first_chunk_size
                current_chunk_step = 0
                active_chunk_latents = None
            except PipelineTimeoutException as e:
                state_to_save = {
                    "accumulated_latents": None,
                    "active_chunk_latents": e.latents.cpu(),
                    "current_chunk_step": e.actual_idx,
                    "current_frame_offset": 0,
                    "total_target_frames": total_target_frames,
                    "prompt": prompt,
                    "resolution": resolution,
                    "seed": seed,
                    "img_bytes": img_bytes,
                    "img_name": img_name,
                    "aud_bytes": aud_bytes,
                    "aud_name": aud_name,
                }
                temp_state_file = os.path.join(tempfile.gettempdir(), f"avatar_state_{uuid.uuid4().hex[:6]}.pt")
                torch.save(state_to_save, temp_state_file)
                status_msg = (
                    f"⏸️ **Processing time for this session has ended.**\n\n"
                    f"📦 **Checkpoint saved:** progress up to frame **0** "
                    f"(of **{total_target_frames}** total frames) has been saved to a `.pt` state file — "
                    f"this is a snapshot of exactly where generation stopped, not an error.\n\n"
                    f"👇 **To continue:**\n"
                    f"1. **Download** the state file from the box on the right.\n"
                    f"2. This file can be uploaded again at any time — in a new session here, on your own "
                    f"local GPU, or anywhere else that runs this pipeline — with no loss of progress.\n"
                    f"3. Upload it into the **\"Upload state file (.pt)\"** field and click Generate again: "
                    f"rendering will resume from the remaining diffusion step."
                )
                yield None, temp_state_file, gr.Markdown(status_msg)
                return
        else:
            # Continue the video autoregressively from the previous cut, using a sliding window
            start_frame = current_frame_offset - 13
            end_frame = start_frame + num_frames_chunk
            audio_emb_chunk = audio_emb[:, start_frame:end_frame]
            
            # If the tail of the audio file is shorter than the chunk's embedding range, pad by repeating the last frame
            if audio_emb_chunk.shape[1] < num_frames_chunk:
                pad_len = num_frames_chunk - audio_emb_chunk.shape[1]
                pad_tensor = audio_emb_chunk[:, -1:].repeat(1, pad_len, 1, 1, 1)
                audio_emb_chunk = torch.cat([audio_emb_chunk, pad_tensor], dim=1)

            if active_chunk_latents is not None:
                cond_latents = None
                latents_input = active_chunk_latents.to(device=device, dtype=dit_dtype)
                ref_latent = accumulated_latents[:, :, :1].to(device=device, dtype=dit_dtype) if accumulated_latents is not None else None
            else:
                ref_latent = accumulated_latents[:, :, :1].to(device=device, dtype=dit_dtype)
                cond_latents = accumulated_latents[:, :, -4:].to(device=device, dtype=dit_dtype)
                latents_input = None

            try:
                with torch.inference_mode():
                    orig_vcond = pipe.use_vcond
                    pipe.use_vcond = True
                    dummy_video = [image] * 13

                    chunk_latents = pipe.generate_avc(
                        video=dummy_video,
                        video_latent=cond_latents,
                        prompt=prompt,
                        negative_prompt=NEGATIVE_PROMPT,
                        height=height,
                        width=width,
                        num_frames=num_frames_chunk,
                        num_cond_frames=13,
                        num_inference_steps=8,
                        text_guidance_scale=1.0,
                        audio_guidance_scale=1.0,
                        output_type="latent",
                        generator=generator,
                        audio_emb=audio_emb_chunk,
                        use_distill=True,
                        use_kv_cache=True,
                        enhance_hf=False,
                        ref_latent=ref_latent,
                        ref_img_index=0,
                        latents=latents_input,
                        start_step=current_chunk_step,
                        start_time=start_gpu_time,
                        timeout=95.0,
                    )
                    pipe.use_vcond = orig_vcond
                
                # Merge the new latents, dropping the exact mathematical overlap (overlap slice starts at index 4)
                new_latents = chunk_latents[:, :, 4:].cpu()
                accumulated_latents = torch.cat([accumulated_latents, new_latents], dim=2)
                current_frame_offset += (num_frames_chunk - 13)
                current_chunk_step = 0
                active_chunk_latents = None
            except PipelineTimeoutException as e:
                state_to_save = {
                    "accumulated_latents": accumulated_latents.cpu(),
                    "active_chunk_latents": e.latents.cpu(),
                    "current_chunk_step": e.actual_idx,
                    "current_frame_offset": current_frame_offset,
                    "total_target_frames": total_target_frames,
                    "prompt": prompt,
                    "resolution": resolution,
                    "seed": seed,
                    "img_bytes": img_bytes,
                    "img_name": img_name,
                    "aud_bytes": aud_bytes,
                    "aud_name": aud_name,
                }
                temp_state_file = os.path.join(tempfile.gettempdir(), f"avatar_state_{uuid.uuid4().hex[:6]}.pt")
                torch.save(state_to_save, temp_state_file)
                status_msg = (
                    f"⏸️ **Processing time for this session has ended.**\n\n"
                    f"📦 **Checkpoint saved:** progress up to frame **{current_frame_offset}** "
                    f"(of **{total_target_frames}** total frames) has been saved to a `.pt` state file — "
                    f"this is a snapshot of exactly where generation stopped, not an error.\n\n"
                    f"👇 **To continue:**\n"
                    f"1. **Download** the state file from the box on the right.\n"
                    f"2. This file can be uploaded again at any time — in a new session here, on your own "
                    f"local GPU, or anywhere else that runs this pipeline — with no loss of progress. "
                    f"Upload it into the **\"Upload state file (.pt)\"** field and click Generate again to "
                    f"resume rendering from the remaining diffusion step."
                )
                yield None, temp_state_file, gr.Markdown(status_msg)
                return

    # 3. Final render of the accumulated latents once video generation is fully complete
    # As with the generation loop, this happens frame-by-frame with safe pause/resume support,
    # so that on long videos, running past the available GPU session time doesn't raise a RuntimeError.
    total_latent_frames = (total_target_frames - 1) // 4 + 1
    final_latents = accumulated_latents[:, :, :total_latent_frames].to(dtype=pipe.vae.dtype)
    final_latents = pipe.denormalize_latents(final_latents)

    t_dec = time.perf_counter()
    while True:
        remaining = total_latent_frames - decode_frame_idx
        yield None, None, gr.Markdown(
            f"🎬 **Final rendering: denoising and compositing the high-quality output (VAE decode)...** "
            f"latent frame **{decode_frame_idx}** of **{total_latent_frames}** (remaining: {remaining})"
        )

        decode_result = _chunked_vae_decode_step(
            pipe, final_latents, decode_frame_idx, decode_feat_map, decode_conv_idx,
            start_gpu_time=start_gpu_time, timeout=95.0,
        )

        # Encode the frames rendered in this session and append them to the accumulated silent video
        if decode_result["frames"]:
            session_video_path = os.path.join(tempfile.gettempdir(), f"decode_part_{uuid.uuid4().hex[:8]}.mp4")
            _encode_frames_silent_mp4(decode_result["frames"], session_video_path, fps=save_fps, quality=5)

            prev_video_path = None
            if decode_video_bytes:
                prev_video_path = os.path.join(tempfile.gettempdir(), f"decode_prev_{uuid.uuid4().hex[:8]}.mp4")
                with open(prev_video_path, "wb") as f:
                    f.write(decode_video_bytes)

            merged_video_path = os.path.join(tempfile.gettempdir(), f"decode_merged_{uuid.uuid4().hex[:8]}.mp4")
            _concat_mp4_stream_copy([prev_video_path, session_video_path], merged_video_path)

            with open(merged_video_path, "rb") as f:
                decode_video_bytes = f.read()

            for p in (session_video_path, prev_video_path, merged_video_path):
                if p and os.path.exists(p):
                    try:
                        os.remove(p)
                    except OSError:
                        pass

        decoded_pixel_frame_count += len(decode_result["frames"])
        decode_frame_idx = decode_result["next_frame_idx"]
        decode_feat_map = decode_result["feat_map"]
        decode_conv_idx = decode_result["conv_idx"]

        del decode_result
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if decode_frame_idx >= total_latent_frames:
            break

        # The session's time budget is running out: save the final-render state (VAE cache + partial video) to resume in the next session
        state_to_save = {
            "phase": "decode",
            "accumulated_latents": accumulated_latents,
            "current_frame_offset": total_target_frames,
            "total_target_frames": total_target_frames,
            "prompt": prompt,
            "resolution": resolution,
            "seed": seed,
            "img_bytes": img_bytes,
            "img_name": img_name,
            "aud_bytes": aud_bytes,
            "aud_name": aud_name,
            "current_chunk_step": 0,
            "active_chunk_latents": None,
            "decode_frame_idx": decode_frame_idx,
            "decode_feat_map": decode_feat_map,
            "decode_conv_idx": decode_conv_idx,
            "decode_video_bytes": decode_video_bytes,
            "decoded_pixel_frame_count": decoded_pixel_frame_count,
        }
        temp_state_file = os.path.join(tempfile.gettempdir(), f"avatar_state_{uuid.uuid4().hex[:6]}.pt")
        torch.save(state_to_save, temp_state_file)
        status_msg = (
            f"⏸️ **Processing time for this session has ended.**\n\n"
            f"📦 **Checkpoint saved:** final rendering (VAE decode) reached frame **{decode_frame_idx}** "
            f"of **{total_latent_frames}** latent frames and has been saved to a `.pt` state file — "
            f"with no skipped or discontinuous frames going forward.\n\n"
            f"👇 **To continue:**\n"
            f"1. **Download** the state file from the box on the right.\n"
            f"2. This file can be uploaded again at any time — in a new session here, on your own "
            f"local GPU, or anywhere else that runs this pipeline — with no loss of progress.\n"
            f"3. Upload it into the **\"Upload state file (.pt)\"** field and click Generate again: "
            f"final rendering will resume from the same frame."
        )
        yield None, temp_state_file, gr.Markdown(status_msg)
        return

    print(f"[timing] vae_decode_total={time.perf_counter() - t_dec:.2f}s", flush=True)

    yield None, None, gr.Markdown("🔊 **Muxing the original audio track with the rendered video...**")
    video_only_path = os.path.join(tempfile.gettempdir(), f"decode_final_{uuid.uuid4().hex[:8]}.mp4")
    with open(video_only_path, "wb") as f:
        f.write(decode_video_bytes)

    out_base = Path(tempfile.gettempdir()) / f"longcat_{uuid.uuid4().hex[:8]}"
    out_path = _mux_final_video(video_only_path, audio_path, str(out_base) + ".mp4", decoded_pixel_frame_count, fps=save_fps)
    print(f"[gen] Successfully completed long video generation. Path: {out_path}", flush=True)

    try:
        os.remove(video_only_path)
    except OSError:
        pass

    # Clean up the temporary uploaded state-file cache
    if state_file and os.path.exists(state_file):
        try:
            os.remove(state_file)
        except OSError:
            pass

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    yield out_path, None, "✅ Video generation completed successfully and audio is fully synced!"


# ---------------------------------------------------------------------------
# 5) Gradio UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="inference-checkpoint-saver", css=CUSTOM_CSS) as demo:
    gr.Markdown(
        """
        # 🎤 LongCat-Video-Avatar 1.5: Audio-Image-to-Video

        ## A reference implementation of checkpointed inference for long-running AI workloads

        Many AI models — video generation, audio generation, and other long-running inference
        workloads — need more processing time than a single available GPU session provides. This
        Space, **inference-checkpoint-saver**, is a reference implementation of a general-purpose
        technique for that situation: **checkpointed save & resume inference**, where a model's
        generation state is written out mid-process and can be reloaded later to resume from the
        exact point it left off, with no loss of progress and no discontinuity in the output.

        We use **LongCat-Video-Avatar-1.5-2nd**, an audio-driven video generation model, as one
        concrete example implementation of this approach. While processing is underway, the full
        generation state — accumulated video latents, the decoder's cache, the current diffusion
        step, and the reference inputs — is serialized to a portable state file. That file can be
        loaded in any later session, and generation continues exactly where it left off, frame by
        frame, with no skipped or discontinuous output. This is what makes it possible to produce up
        to 2 minutes of audio-driven video from a single reference image, well beyond what one
        continuous processing session would otherwise allow.

        This pattern generalizes to any autoregressive or chunked generation model that needs to
        produce long-form output under a limited amount of continuous compute time, not just this one.

        Upload a reference image, a driving audio clip, and a short text prompt to get started.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            image_in = gr.Image(label="Reference image", type="filepath")
            audio_in = gr.Audio(label="Driving audio", type="filepath")
            prompt = gr.Textbox(
                label="Prompt",
                value="A person is speaking expressively, looking at the camera.",
                lines=3,
            )
            resolution = gr.State(value="480p")
            seed = gr.Number(value=42, precision=0, label="Seed")
            vocal_mode = gr.Radio(
                [VOCAL_MODE_FAST, VOCAL_MODE_QUALITY],
                value=VOCAL_MODE_FAST,
                label="Audio preprocessing",
            )
            acceleration = gr.Radio(
                [ACCEL_MODE_EXACT, ACCEL_MODE_DBCACHE, ACCEL_MODE_DBCACHE_FASTER],
                value=ACCEL_MODE_DBCACHE_FASTER,
                label="Acceleration",
            )
            tracking_id_in = gr.File(
                label="Upload state file (.pt)",
                file_types=[".pt"],
                type="filepath"
            )
            go = gr.Button("Generate", variant="primary")
        with gr.Column(scale=1):
            video_out = gr.Video(label="Output Video", autoplay=True, height=420)
            tracking_id_out = gr.File(
                label="New state file (.pt) — download this and re-upload it above if generation times out",
                type="filepath",
                interactive=False
            )
            move_state_up_btn = gr.Button("⬆️ Move state file up (no download/re-upload needed)")
            status_out = gr.Markdown(value="**Status:** waiting to start...", label="System Status")

    def _move_state_file_up(state_path):
        # The file already exists on the same server; only its path is moved into the upload box,
        # so there's no need to download it in the browser and re-upload it (saving the user's bandwidth).
        if not state_path:
            return None, gr.Markdown("⚠️ No state file has been generated yet to move.")
        return state_path, gr.Markdown("✅ State file moved into the \"Upload state file\" field. Click Generate to continue.")

    move_state_up_btn.click(
        _move_state_file_up,
        inputs=[tracking_id_out],
        outputs=[tracking_id_in, status_out],
    )

    go.click(
        generate,
        inputs=[image_in, audio_in, prompt, resolution, seed, vocal_mode, acceleration, tracking_id_in],
        outputs=[video_out, tracking_id_out, status_out],
    )

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