app.py

import spaces
import os
import gc
import cv2
import json
import uuid
import shutil
import subprocess
import tempfile
import threading
from pathlib import Path

import torch
print("Torch:", torch.__version__)
print("Torch CUDA:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())

import numpy as np
import gradio as gr

from PIL import Image
from huggingface_hub import snapshot_download, hf_hub_download

# ============================================================
# DIFFUSERS
# ============================================================

import diffusers

print("Diffusers version:", diffusers.__version__)
print(
    "Has WanVACEPipeline:",
    hasattr(diffusers, "WanVACEPipeline")
)

from diffusers import (
    WanVACEPipeline,
    WanVACETransformer3DModel,
    AutoencoderKLWan,
)



from diffusers.schedulers.scheduling_unipc_multistep import (
    UniPCMultistepScheduler,
)

from diffusers.utils import export_to_video

# ============================================================
# SAM 3
# ============================================================

try:
    from sam3.model_builder import build_sam3_video_predictor
except Exception as _sam3_import_error:
    build_sam3_video_predictor = None
    _SAM3_IMPORT_ERROR = _sam3_import_error
else:
    _SAM3_IMPORT_ERROR = None


# ============================================================
# CONFIG
# ============================================================

WAN_MODEL = os.environ.get(
    "WAN_MODEL",
    "Wan-AI/Wan2.1-VACE-14B-diffusers",
)

# SAM 3.0 is used for automatic text-prompted segmentation/tracking.
# SAM 3.1 Object Multiplex is intentionally not used here because the public
# checkpoint/code pair has had compatibility issues in current releases.
SAM3_MODEL = os.environ.get(
    "SAM3_MODEL",
    "facebook/sam3",
)
SAM3_CHECKPOINT_NAME = os.environ.get(
    "SAM3_CHECKPOINT_NAME",
    "sam3.pt",
)

# Prefer persistent HF Space storage when it exists; otherwise fall back
# to /tmp for a single-runtime cache.
_DEFAULT_MODEL_ROOT = (
    Path("/data/zero_drift_models")
    if Path("/data").exists()
    else Path("/tmp/zero_drift_models")
)

MODEL_ROOT = Path(
    os.environ.get(
        "MODEL_ROOT",
        str(_DEFAULT_MODEL_ROOT),
    )
)

WAN_LOCAL_DIR = MODEL_ROOT / "wan_vace"
SAM3_LOCAL_DIR = MODEL_ROOT / "sam3"

SAM3_CHECKPOINT = SAM3_LOCAL_DIR / SAM3_CHECKPOINT_NAME
SAM3_BPE = SAM3_LOCAL_DIR / "bpe_simple_vocab_16e6.txt.gz"

# ZeroGPU exposes CUDA when the @spaces.GPU function is executing.
DEVICE = "cuda"

WAN_DTYPE = torch.bfloat16

# Favor Tensor Core execution for the remaining FP32 helper operations.
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass
try:
    torch.backends.cuda.matmul.allow_tf32 = True
except Exception:
    pass
try:
    torch.backends.cudnn.allow_tf32 = True
except Exception:
    pass

# Full-quality 14B BF16 path. The 96 GB xlarge GPU has enough VRAM,
# so do not quantize the transformer: native BF16 Tensor Core execution
# is the latency-first path here.
USE_TORCH_COMPILE = False  # Stable eager path; avoids the current Dynamo/CUDA handler bug.

OUTPUT_DIR = Path(
    "/tmp/zero_drift_outputs"
)

OUTPUT_DIR.mkdir(
    parents=True,
    exist_ok=True,
)

# Maximum source duration.
MAX_SECONDS = 15

# Wan 2.1 VACE is trained around short 4k+1-frame clips.
# Longer source videos are processed in overlapping chunks.
VACE_CHUNK_FRAMES = 81
VACE_CHUNK_OVERLAP = 8
MODEL_FPS = 16.0

# 14B quality path: keep the full 720p-capable spatial budget.
# Speed is prioritized with native BF16 execution and low step counts.
MAX_WIDTH = 1280
MAX_HEIGHT = 720


# ============================================================
# GLOBAL MODELS
# ============================================================

WAN_PIPE = None
SAM3_PREDICTOR = None
MODEL_LOCK = threading.Lock()


# ============================================================
# STARTUP MODEL DOWNLOADS
# ============================================================

def ensure_model_downloads():
    """
    Download every required model artifact exactly once per runtime.

    Wan is downloaded as a complete local Diffusers repository.
    SAM 3.0 is downloaded from Hugging Face.
    """

    MODEL_ROOT.mkdir(
        parents=True,
        exist_ok=True,
    )

    WAN_LOCAL_DIR.mkdir(
        parents=True,
        exist_ok=True,
    )

    SAM3_LOCAL_DIR.mkdir(
        parents=True,
        exist_ok=True,
    )

    print("=" * 72)
    print("ZERO-DRIFT STARTUP MODEL DOWNLOAD")
    print("Model root:", MODEL_ROOT)
    print("Wan repo:", WAN_MODEL)
    print("SAM 3 repo:", SAM3_MODEL)
    print("=" * 72)

    # --------------------------------------------------------
    # WAN 2.1
    # --------------------------------------------------------

    wan_config = WAN_LOCAL_DIR / "model_index.json"
    wan_vae_config = WAN_LOCAL_DIR / "vae" / "config.json"
    wan_vae_weights = WAN_LOCAL_DIR / "vae" / "diffusion_pytorch_model.safetensors"
    wan_marker = WAN_LOCAL_DIR / ".download_complete"

    wan_cache_valid = (
        wan_config.exists()
        and wan_vae_config.exists()
        and wan_vae_weights.exists()
        and wan_marker.exists()
    )

    if wan_cache_valid:
        print("Wan 2.1 VACE Diffusers repo already present:")
        print(" ", WAN_LOCAL_DIR)
    else:
        if WAN_LOCAL_DIR.exists():
            print("Existing Wan cache is not a valid Diffusers VACE repo.")
            print("Removing stale cache:", WAN_LOCAL_DIR)
            shutil.rmtree(WAN_LOCAL_DIR, ignore_errors=True)

        WAN_LOCAL_DIR.mkdir(
            parents=True,
            exist_ok=True,
        )

        print("Downloading Wan 2.1 VACE Diffusers to local storage...")
        print("This repository is about 75 GB.")
        print("It will NOT be downloaded again during this runtime.")

        snapshot_download(
            repo_id=WAN_MODEL,
            local_dir=str(WAN_LOCAL_DIR),
        )

        # Do not mark the download complete until the Diffusers VAE
        # files that AutoencoderKLWan.from_pretrained() requires are present.
        if not wan_vae_config.exists() or not wan_vae_weights.exists():
            raise RuntimeError(
                "Wan VACE Diffusers download completed, but the local "
                "vae/config.json or vae/diffusion_pytorch_model.safetensors "
                "is missing. Check the model repository and download."
            )

        wan_marker.touch()
        print("Wan 2.1 VACE Diffusers download complete.")

    # --------------------------------------------------------
    # SAM 3.0
    # --------------------------------------------------------

    SAM3_LOCAL_DIR.mkdir(
        parents=True,
        exist_ok=True,
    )

    if (
        SAM3_CHECKPOINT.exists()
        and SAM3_CHECKPOINT.stat().st_size > 0
    ):
        print("SAM 3 checkpoint already present:")
        print(" ", SAM3_CHECKPOINT)
    else:
        print("Downloading SAM 3 checkpoint from Hugging Face...")
        print("Repo:", SAM3_MODEL)
        print("File:", SAM3_CHECKPOINT_NAME)
        token = os.environ.get("HF_TOKEN") or None
        hf_hub_download(
            repo_id=SAM3_MODEL,
            filename=SAM3_CHECKPOINT_NAME,
            local_dir=str(SAM3_LOCAL_DIR),
            token=token,
        )

    if not SAM3_CHECKPOINT.exists() or SAM3_CHECKPOINT.stat().st_size == 0:
        raise RuntimeError(
            f"SAM 3 checkpoint missing after download: {SAM3_CHECKPOINT}"
        )

    if not SAM3_BPE.exists():
        print("SAM 3 BPE asset missing; downloading it...")
        import urllib.request
        bpe_url = (
            "https://raw.githubusercontent.com/openai/CLIP/main/clip/"
            "bpe_simple_vocab_16e6.txt.gz"
        )
        urllib.request.urlretrieve(
            bpe_url,
            str(SAM3_BPE),
        )

    print("SAM 3 checkpoint:", SAM3_CHECKPOINT)
    print("SAM 3 BPE:", SAM3_BPE)
    print("=" * 72)


# ============================================================
# GENERAL HELPERS
# ============================================================

def cleanup():
    gc.collect()

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

        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass


def run_command(command):

    process = subprocess.run(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )

    if process.returncode != 0:
        raise RuntimeError(
            "\n".join([
                "Command failed:",
                " ".join(command),
                "",
                process.stderr[-6000:],
            ])
        )

    return process.stdout


def get_video_path(value):

    if value is None:
        return None

    if isinstance(value, str):
        return value

    if isinstance(value, dict):

        if value.get("path"):
            return value["path"]

        if value.get("video"):
            return value["video"]

    raise ValueError(
        f"Unsupported video input: {type(value)}"
    )


# ============================================================
# VIDEO PROBING
# ============================================================

def probe_video(path):

    result = run_command([
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "stream=index,codec_type,width,height,"
        "r_frame_rate,avg_frame_rate,duration",
        "-of",
        "json",
        path,
    ])

    data = json.loads(result)

    video_stream = None
    audio_stream = None

    for stream in data.get("streams", []):

        if stream.get("codec_type") == "video":
            video_stream = stream

        elif stream.get("codec_type") == "audio":
            audio_stream = stream

    if video_stream is None:
        raise RuntimeError(
            "Input does not contain a video stream."
        )

    width = int(
        video_stream["width"]
    )

    height = int(
        video_stream["height"]
    )

    fps_string = (
        video_stream.get("avg_frame_rate")
        or video_stream.get("r_frame_rate")
        or "16/1"
    )

    try:

        n, d = fps_string.split("/")

        fps = float(n) / float(d)

    except Exception:

        fps = 16.0

    duration = float(
        video_stream.get(
            "duration",
            0,
        )
        or 0
    )

    return {
        "width": width,
        "height": height,
        "fps": fps,
        "duration": duration,
        "has_audio": audio_stream is not None,
    }


# ============================================================
# AUDIO
# ============================================================

def extract_audio(
    input_video,
    output_audio,
):

    # -vn means absolutely no video processing.
    # -c:a copy means the original audio stream is copied.

    run_command([
        "ffmpeg",
        "-y",
        "-i",
        input_video,
        "-vn",
        "-map",
        "0:a:0?",
        "-c:a",
        "copy",
        output_audio,
    ])

    return output_audio


# ============================================================
# VIDEO DECODING
# ============================================================

def decode_frames(
    video_path,
    max_frames=None,
    target_fps=MODEL_FPS,
):

    info = probe_video(
        video_path
    )

    cap = cv2.VideoCapture(
        video_path
    )

    if not cap.isOpened():
        raise RuntimeError(
            "Could not open input video."
        )

    source_fps = float(
        info.get("fps") or target_fps
    )

    duration = float(
        info.get("duration") or 0.0
    )

    if target_fps and duration > 0:
        target_count = max(1, int(round(duration * target_fps)))
    else:
        target_count = int(
            cap.get(cv2.CAP_PROP_FRAME_COUNT)
        )

    source_count = max(
        1,
        int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
    )

    if max_frames is not None:
        target_count = min(
            target_count,
            int(max_frames),
        )

    target_count = min(
        target_count,
        source_count,
    )

    target_indices = np.linspace(
        0,
        max(0, source_count - 1),
        target_count,
        dtype=np.int64,
    )

    target_set = {
        int(i): pos
        for pos, i in enumerate(target_indices)
    }

    frames = []
    source_index = 0

    while len(frames) < target_count:

        ok, frame = cap.read()

        if not ok:
            break

        if source_index in target_set:
            frame = cv2.cvtColor(
                frame,
                cv2.COLOR_BGR2RGB,
            )

            frames.append(
                Image.fromarray(frame)
            )

        source_index += 1

    cap.release()

    if not frames:
        raise RuntimeError(
            "Could not decode any frames."
        )

    # The Wan VACE model is used at a stable 16 fps working rate.
    # This preserves the source duration while avoiding unnecessary
    # extra VACE chunks for 24/30/60 fps inputs.
    return (
        frames,
        float(target_fps or source_fps),
        info,
    )


# ============================================================
# RESOLUTION
# ============================================================

def resize_for_wan(
    frames,
):

    width, height = frames[0].size

    scale = min(
        MAX_WIDTH / width,
        MAX_HEIGHT / height,
        1.0,
    )

    new_width = int(
        width * scale
    )

    new_height = int(
        height * scale
    )

    # Wan dimensions should be divisible by 8.
    new_width -= (
        new_width % 8
    )

    new_height -= (
        new_height % 8
    )

    new_width = max(
        256,
        new_width,
    )

    new_height = max(
        256,
        new_height,
    )

    if (
        new_width == width
        and new_height == height
    ):
        return frames

    return [
        frame.resize(
            (
                new_width,
                new_height,
            ),
            Image.Resampling.LANCZOS,
        )
        for frame in frames
    ]


# ============================================================
# MODEL LOADING
# ============================================================

def load_models():
    global WAN_PIPE
    global SAM3_PREDICTOR

    with MODEL_LOCK:

        if WAN_PIPE is None:

            print("Loading Wan 2.1 VACE Diffusers from local files...")
            print("Local path:", WAN_LOCAL_DIR)

            vae_config = WAN_LOCAL_DIR / "vae" / "config.json"
            vae_weights = (
                WAN_LOCAL_DIR
                / "vae"
                / "diffusion_pytorch_model.safetensors"
            )

            if not vae_config.exists() or not vae_weights.exists():
                raise RuntimeError(
                    "Wan VACE Diffusers cache is incomplete: expected "
                    f"{vae_config} and {vae_weights}."
                )

            # VAE stays FP32 for numerical stability.
            vae = AutoencoderKLWan.from_pretrained(
                str(WAN_LOCAL_DIR),
                subfolder="vae",
                torch_dtype=torch.float32,
                local_files_only=True,
            )

            WAN_PIPE = WanVACEPipeline.from_pretrained(
                str(WAN_LOCAL_DIR),
                vae=vae,
                torch_dtype=WAN_DTYPE,
                local_files_only=True,
            )

            WAN_PIPE.scheduler = UniPCMultistepScheduler.from_config(
                WAN_PIPE.scheduler.config,
                flow_shift=5.0,
            )

            # Full native BF16 on the Blackwell xlarge GPU.
            WAN_PIPE.to(DEVICE)

            # Fused QKV reduces projection launches. This is experimental in
            # Diffusers, so keep it guarded rather than making startup fragile.
            try:
                WAN_PIPE.fuse_qkv_projections()
                print("VACE QKV projections fused.")
            except Exception as e:
                print("QKV fusion skipped:", repr(e))

            # Do not use torch.compile on this ZeroGPU runtime. PyTorch 2.13 +
            # this CUDA build can fail lazily on the first transformer call with
            # an internal Dynamo handler-registration AssertionError. Native
            # BF16 eager execution is therefore the stable path.
            print("torch.compile disabled; using native BF16 eager execution.")

            WAN_PIPE.set_progress_bar_config(disable=False)
            print("Wan VACE 14B BF16 loaded.")

        if SAM3_PREDICTOR is None:

            if build_sam3_video_predictor is None:
                raise RuntimeError(
                    "Official SAM 3 is not installed. Install it with "
                    "git+https://github.com/facebookresearch/sam3.git. "
                    f"Original import error: {_SAM3_IMPORT_ERROR!r}"
                )

            print("Loading SAM 3.0 video predictor from local files...")

            SAM3_PREDICTOR = build_sam3_video_predictor(
                checkpoint_path=str(SAM3_CHECKPOINT),
                bpe_path=str(SAM3_BPE),
                gpus_to_use=[0],
                compile=False,
            )

            print("SAM 3 loaded.")


# ============================================================
# SAM 3 AUTOMATIC TEXT SEGMENTATION + TRACKING
# ============================================================

def _sam3_masks_from_outputs(outputs):
    if outputs is None:
        return [], []
    obj_ids = outputs.get("out_obj_ids")
    masks = outputs.get("out_binary_masks")
    if obj_ids is None or masks is None:
        raise RuntimeError(
            "SAM 3 returned no object masks. "
            f"Available keys: {list(outputs.keys())}"
        )
    if torch.is_tensor(obj_ids):
        obj_ids = obj_ids.detach().cpu().numpy()
    else:
        obj_ids = np.asarray(obj_ids)
    if torch.is_tensor(masks):
        masks = masks.detach().cpu().numpy()
    else:
        masks = np.asarray(masks)
    if masks.ndim == 2:
        masks = masks[None, ...]
    return obj_ids.astype(int).tolist(), list(masks.astype(bool))


def track_object_text(video_path, target_text, expected_frames):
    if not target_text or not target_text.strip():
        raise RuntimeError(
            "Enter what should be segmented, e.g. 'upper torso', 'shirt', "
            "'white bikini', or 'person'."
        )
    if SAM3_PREDICTOR is None:
        raise RuntimeError("SAM 3 predictor is not loaded.")

    response = SAM3_PREDICTOR.handle_request({
        "type": "start_session",
        "resource_path": str(video_path),
        "offload_video_to_cpu": True,
        "offload_state_to_cpu": True,
    })
    session_id = response["session_id"]
    try:
        print(f"SAM 3 target: {target_text!r}")
        prompted = SAM3_PREDICTOR.handle_request({
            "type": "add_prompt",
            "session_id": session_id,
            "frame_index": 0,
            "text": target_text.strip(),
            "output_prob_thresh": 0.40,
        })
        first_ids, first_masks = _sam3_masks_from_outputs(prompted.get("outputs"))
        if not first_masks:
            raise RuntimeError(
                f"SAM 3 found no object matching '{target_text}'. Try a more concrete target such as 'upper torso', 'shirt', or 'bikini'."
            )

        areas = [int(m.sum()) for m in first_masks]
        chosen_index = int(np.argmax(areas))
        chosen_obj_id = int(first_ids[chosen_index])

        tracked = {0: first_masks[chosen_index].astype(np.uint8)}

        stream = SAM3_PREDICTOR.handle_stream_request({
            "type": "propagate_in_video",
            "session_id": session_id,
            "propagation_direction": "forward",
            "start_frame_index": 0,
            "max_frame_num_to_track": int(expected_frames),
            "output_prob_thresh": 0.40,
        })

        for item in stream:
            frame_idx = int(item["frame_index"])
            obj_ids, masks = _sam3_masks_from_outputs(item.get("outputs"))
            for obj_id, mask in zip(obj_ids, masks):
                if int(obj_id) == chosen_obj_id:
                    tracked[frame_idx] = mask.astype(np.uint8)
                    break

        height, width = first_masks[0].shape[-2:]
        output = []
        previous = tracked.get(0, np.zeros((height, width), dtype=np.uint8))
        for i in range(int(expected_frames)):
            mask = tracked.get(i, previous)
            if mask.shape != (height, width):
                mask = cv2.resize(mask.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST)
            previous = mask
            output.append(mask)
        return output
    finally:
        try:
            SAM3_PREDICTOR.handle_request({
                "type": "close_session",
                "session_id": session_id,
                "run_gc_collect": True,
                "clear_cache_threshold": 90,
            })
        except Exception as e:
            print("SAM 3 session close warning:", repr(e))


# ============================================================
# FACE PROTECTION
# ============================================================

FACE_CASCADE = cv2.CascadeClassifier(
    cv2.data.haarcascades
    + "haarcascade_frontalface_default.xml"
)


def face_mask(frame):

    rgb = np.asarray(
        frame.convert("RGB")
    )

    gray = cv2.cvtColor(
        rgb,
        cv2.COLOR_RGB2GRAY,
    )

    faces = FACE_CASCADE.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(40, 40),
    )

    mask = np.zeros(
        gray.shape,
        dtype=np.uint8,
    )

    for (
        x,
        y,
        w,
        h,
    ) in faces:

        pad_x = int(
            w * 0.20
        )

        pad_y = int(
            h * 0.25
        )

        x1 = max(
            0,
            x - pad_x,
        )

        y1 = max(
            0,
            y - pad_y,
        )

        x2 = min(
            gray.shape[1],
            x + w + pad_x,
        )

        y2 = min(
            gray.shape[0],
            y + h + pad_y,
        )

        mask[
            y1:y2,
            x1:x2,
        ] = 1

    # Additional dilation.
    kernel = np.ones(
        (11, 11),
        dtype=np.uint8,
    )

    mask = cv2.dilate(
        mask,
        kernel,
        iterations=1,
    )

    return mask


# ============================================================
# FINAL EDIT MASK
# ============================================================

def create_edit_masks(
    target_masks,
    frames,
):

    output = []

    for target, frame in zip(
        target_masks,
        frames,
    ):

        protected = face_mask(
            frame
        )

        target = target.astype(
            np.uint8
        )

        # Target minus protected face.
        editable = target * (
            1 - protected
        )

        # Feather boundary.
        editable = cv2.GaussianBlur(
            editable.astype(
                np.float32
            ),
            (0, 0),
            sigmaX=3.0,
        )

        editable = np.clip(
            editable,
            0.0,
            1.0,
        )

        output.append(
            editable
        )

    return output


# ============================================================
# PIXEL COMPOSITING
# ============================================================

def _frame_to_rgb_array(frame):
    """
    Normalize PIL / NumPy / torch outputs into HxWx3 uint8.
    Wan VACE may return NumPy arrays, while other versions can return PIL.
    """
    if isinstance(frame, Image.Image):
        arr = np.asarray(frame.convert("RGB"))
    elif torch.is_tensor(frame):
        arr = frame.detach().float().cpu().numpy()
    else:
        arr = np.asarray(frame)

    if arr.ndim == 4 and arr.shape[0] == 1:
        arr = arr[0]

    # Handle CHW tensors/arrays.
    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
        arr = np.transpose(arr, (1, 2, 0))

    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)

    if arr.shape[-1] == 4:
        arr = arr[..., :3]

    if np.issubdtype(arr.dtype, np.floating):
        finite_max = float(np.nanmax(arr)) if arr.size else 0.0
        if finite_max <= 1.5:
            arr = arr * 255.0

    return np.clip(arr, 0, 255).astype(np.uint8)


def composite(
    original,
    generated,
    mask,
):
    original = _frame_to_rgb_array(original).astype(np.float32)
    generated = _frame_to_rgb_array(generated).astype(np.float32)

    if generated.shape != original.shape:
        generated = cv2.resize(
            generated,
            (
                original.shape[1],
                original.shape[0],
            ),
            interpolation=cv2.INTER_LANCZOS4,
        )

    mask = cv2.resize(
        mask.astype(np.float32),
        (
            original.shape[1],
            original.shape[0],
        ),
        interpolation=cv2.INTER_LINEAR,
    )

    mask = np.clip(mask, 0.0, 1.0)[..., None]

    result = (
        generated * mask
        + original * (1.0 - mask)
    )

    return Image.fromarray(
        np.clip(
            result,
            0,
            255,
        ).astype(np.uint8)
    )


# ============================================================
# WAN GENERATION
# ============================================================

def generate_wan(
    frames,
    prompt,
    edit_masks,
    reference,
    strength,
    steps,
    guidance,
    seed,
):

    generator = torch.Generator(
        device=DEVICE
    ).manual_seed(
        int(seed)
    )

    # Wan VACE is the native Wan 2.1 pipeline for masked
    # video-to-video editing plus optional reference-image
    # conditioning. Black mask pixels are kept as conditioning;
    # white pixels are generated/edited.
    vace_masks = [
        Image.fromarray(
            np.clip(
                mask * 255.0,
                0,
                255,
            ).astype(np.uint8),
            mode="L",
        )
        for mask in edit_masks
    ]

    kwargs = dict(
        video=frames,
        mask=vace_masks,
        prompt=prompt,
        negative_prompt=(
            "low quality, worst quality, blurry, "
            "distorted face, deformed face, "
            "bad anatomy, extra limbs, "
            "warped hands, flickering, "
            "frame artifacts, subtitles, "
            "watermark, text"
        ),
        height=frames[0].height,
        width=frames[0].width,
        num_frames=len(frames),
        num_inference_steps=int(
            steps
        ),
        guidance_scale=float(
            guidance
        ),
        conditioning_scale=float(strength),
        generator=generator,
    )

    if reference is not None:
        kwargs["reference_images"] = [
            reference
        ]

    # VACE uses the explicit edit mask rather than the V2V strength
    # parameter exposed by WanVideoToVideoPipeline. The UI strength
    # is applied later during deterministic pixel compositing.
    with torch.inference_mode():
        result = WAN_PIPE(
            **kwargs
        )

    return result.frames[0]


# ============================================================
# VACE CHUNKING
# ============================================================

def _vace_model_length(frame_count):
    """Round a requested chunk up to Wan's 4n+1 temporal shape."""
    frame_count = max(5, int(frame_count))
    return ((frame_count - 1 + 3) // 4) * 4 + 1


def build_vace_chunks(frames, masks):
    """Yield overlapping source chunks without dropping video duration."""
    total = len(frames)
    step = VACE_CHUNK_FRAMES - VACE_CHUNK_OVERLAP

    if total <= VACE_CHUNK_FRAMES:
        yield 0, frames, masks, total
        return

    start = 0
    while start < total:
        end = min(start + VACE_CHUNK_FRAMES, total)
        desired = end - start
        model_len = _vace_model_length(desired)

        chunk_frames = list(frames[start:end])
        chunk_masks = list(masks[start:end])

        # Pad only the short final chunk. These padding frames are
        # discarded again after generation, so the output frame count
        # exactly matches the source working-rate frame count.
        while len(chunk_frames) < model_len:
            chunk_frames.append(chunk_frames[-1].copy())
            chunk_masks.append(chunk_masks[-1].copy())

        yield start, chunk_frames, chunk_masks, desired

        if end >= total:
            break

        start += step


# ============================================================
# ENCODE
# ============================================================

def encode_video(
    frames,
    fps,
    output,
):

    # Diffusers export helper.
    export_to_video(
        frames,
        str(output),
        fps=fps,
    )

    return str(output)


# ============================================================
# AUDIO REMUX
# ============================================================

def remux_audio(
    silent_video,
    original_video,
    output,
):

    # We take the ORIGINAL audio directly.
    #
    # Video is re-encoded only once here.
    # Audio remains untouched.

    run_command([
        "ffmpeg",
        "-y",
        "-i",
        silent_video,
        "-i",
        original_video,
        "-map",
        "0:v:0",
        "-map",
        "1:a:0?",
        "-c:v",
        "libx264",
        "-preset",
        "medium",
        "-crf",
        "16",
        "-pix_fmt",
        "yuv420p",
        "-c:a",
        "copy",
        "-shortest",
        output,
    ])

    return str(output)


# ============================================================
# MAIN GPU FUNCTION
# ============================================================

@spaces.GPU(
    size="xlarge",
    # Override with GPU_DURATION if your Space quota permits a longer job.
    duration=int(os.environ.get("GPU_DURATION", "360")),
)
def process_video(
    video,
    reference,
    target_text,
    prompt,
    strength,
    steps,
    guidance,
    seed,
):

    video_path = get_video_path(
        video
    )

    if video_path is None:
        raise gr.Error(
            "Upload a video."
        )

    if not prompt:
        raise gr.Error(
            "Enter an editing prompt."
        )

    try:

        # ----------------------------------------------------
        # MODEL INITIALIZATION
        # ----------------------------------------------------

        load_models()

        # ----------------------------------------------------
        # INPUT VALIDATION
        # ----------------------------------------------------

        info = probe_video(
            video_path
        )

        if (
            info["duration"]
            > MAX_SECONDS
        ):

            raise gr.Error(
                f"Maximum video length is "
                f"{MAX_SECONDS} seconds."
            )

        # ----------------------------------------------------
        # WORKSPACE
        # ----------------------------------------------------

        job_id = uuid.uuid4().hex

        work = (
            OUTPUT_DIR
            / job_id
        )

        work.mkdir(
            parents=True,
            exist_ok=True,
        )

        # ----------------------------------------------------
        # AUDIO
        # ----------------------------------------------------

        audio_path = (
            work / "audio.mka"
        )

        if info["has_audio"]:

            extract_audio(
                video_path,
                str(audio_path),
            )

        # ----------------------------------------------------
        # FRAMES
        # ----------------------------------------------------

        frames, fps, _ = decode_frames(
            video_path
        )

        frames = resize_for_wan(
            frames
        )

        # ----------------------------------------------------
        # SOURCE FRAME COUNT
        # ----------------------------------------------------

        # Do not truncate a longer source to 81 frames. VACE is
        # processed in overlapping <=81-frame chunks so the final
        # video retains the full source duration.
        if len(frames) < 5:
            raise gr.Error(
                "Video is too short for Wan VACE."
            )

        # ----------------------------------------------------
        # SAM 3 AUTO TARGET SEGMENTATION
        # ----------------------------------------------------

        if not target_text:
            raise gr.Error(
                "Enter what you want to edit, e.g. 'upper torso', 'shirt', 'white bikini', or 'person'."
            )

        sam_video = work / "sam_working.mp4"
        encode_video(frames, MODEL_FPS, sam_video)

        print("Running SAM 3 automatic text segmentation...")
        tracked_masks = track_object_text(
            sam_video,
            target_text,
            len(frames),
        )

        # ----------------------------------------------------
        # FACE PROTECTION
        # ----------------------------------------------------

        print(
            "Building protected regions..."
        )

        edit_masks = (
            create_edit_masks(
                tracked_masks,
                frames,
            )
        )

        # ----------------------------------------------------
        # PROMPT
        # ----------------------------------------------------

        final_prompt = prompt

        if reference is not None:

            final_prompt += (
                ". Use the supplied reference image "
                "as a visual reference for the "
                "selected region."
            )

        # ----------------------------------------------------
        # WAN VACE — FULL VIDEO, CHUNKED
        # ----------------------------------------------------

        print(
            f"Running Wan VACE 14B in <= {VACE_CHUNK_FRAMES}-frame chunks..."
        )

        final_frames = []
        first_chunk = True

        for chunk_index, (
            start,
            chunk_frames,
            chunk_masks,
            desired_count,
        ) in enumerate(
            build_vace_chunks(frames, edit_masks)
        ):

            print(
                f"VACE chunk {chunk_index + 1}: "
                f"source frames {start}..{start + desired_count - 1} "
                f"(model length={len(chunk_frames)})"
            )

            generated = generate_wan(
                frames=chunk_frames,
                prompt=final_prompt,
                edit_masks=chunk_masks,
                reference=reference,
                strength=float(strength),
                steps=int(steps),
                guidance=float(guidance),
                seed=int(seed),
            )

            generated = list(generated[:desired_count])

            # Preserve the exact source frame count. Overlap frames are
            # used as temporal context for VACE but emitted only once.
            if first_chunk:
                emit_from = 0
                first_chunk = False
            else:
                emit_from = min(
                    VACE_CHUNK_OVERLAP,
                    len(generated),
                )

            for local_idx in range(emit_from, len(generated)):
                global_idx = start + local_idx

                # The first chunk owns frames [0, chunk_end). For later
                # chunks, the overlap belongs to the previous chunk.
                if global_idx >= len(frames):
                    break

                final_frames.append(
                    composite(
                        frames[global_idx],
                        generated[local_idx],
                        np.clip(
                            chunk_masks[local_idx],
                            0.0,
                            1.0,
                        ),
                    )
                )

        if len(final_frames) != len(frames):
            raise RuntimeError(
                f"Internal frame accounting error: "
                f"expected {len(frames)} frames, got {len(final_frames)}."
            )

        # ----------------------------------------------------
        # VIDEO
        # ----------------------------------------------------

        silent_video = (
            work / "video.mp4"
        )

        encode_video(
            final_frames,
            fps,
            silent_video,
        )

        # ----------------------------------------------------
        # AUDIO
        # ----------------------------------------------------

        final_video = (
            work / "final.mp4"
        )

        if info["has_audio"]:

            remux_audio(
                str(silent_video),
                video_path,
                str(final_video),
            )

        else:

            shutil.copy2(
                silent_video,
                final_video,
            )

        print(
            "Done:",
            final_video,
        )

        return str(
            final_video
        )

    except gr.Error:
        raise

    except Exception as e:

        print(
            "PIPELINE ERROR:",
            repr(e),
        )

        raise gr.Error(
            f"{type(e).__name__}: {e}"
        )

    finally:

        cleanup()


# ============================================================
# DOWNLOAD MODELS BEFORE THE APP STARTS
# ============================================================

# This is deliberately outside @spaces.GPU.
# Network/model download does not need a GPU and happens once at startup.
ensure_model_downloads()


# ============================================================
# UI
# ============================================================

CSS = """
#title {
    text-align: center;
    font-size: 34px;
    font-weight: 700;
}

#subtitle {
    text-align: center;
    opacity: 0.7;
}

.gradio-container {
    max-width: 1400px !important;
}
"""


with gr.Blocks(
    title="Zero-Drift VACE Auto Editor"
) as demo:

    gr.Markdown(
        "# Zero-Drift VACE Auto Editor",
        elem_id="title",
    )

    gr.Markdown(
        "SAM 3 automatic text segmentation + Wan 2.1 VACE 14B BF16 + "
        "full-duration chunked editing + protected-region compositing",
        elem_id="subtitle",
    )

    with gr.Row():

        # ====================================================
        # LEFT
        # ====================================================

        with gr.Column(
            scale=1
        ):

            video_input = gr.Video(
                label="Source Video"
            )

            reference_input = gr.Image(
                label="Reference Image (optional)",
                type="pil",
            )

            target_input = gr.Textbox(
                label="What should be edited?",
                value="upper torso",
                placeholder="Examples: upper torso, shirt, white bikini, jacket, person's hair",
                lines=2,
            )

            prompt_input = gr.Textbox(
                label="Edit Instruction",
                placeholder=(
                    "Examples: Add the black bikini shown in the reference image. "
                    "Preserve the person's face, body shape, pose, skin, lighting and background."
                ),
                lines=4,
            )

            strength_input = gr.Slider(
                minimum=0.50,
                maximum=1.50,
                value=1.00,
                step=0.05,
                label="VACE Control Strength",
                info="Controls how strongly the mask/control video constrains VACE. The final composite always fully replaces the selected region.",
            )

            steps_input = gr.Slider(
                minimum=6,
                maximum=20,
                value=10,
                step=1,
                label="Inference Steps (14B BF16 — fastest)",
            )

            guidance_input = gr.Slider(
                minimum=1.0,
                maximum=8.0,
                value=5.0,
                step=0.5,
                label="Guidance",
            )

            seed_input = gr.Number(
                value=42,
                precision=0,
                label="Seed",
            )

            generate_button = gr.Button(
                "GENERATE",
                variant="primary",
            )

        # ====================================================
        # RIGHT
        # ====================================================

        with gr.Column(
            scale=1
        ):

            gr.Markdown(
                "### Automatic segmentation\n\n"
                "SAM 3 finds the requested object/region from text and tracks it through the entire video. "
                "No point selection is required.\n\n"
                "For a bare upper-body source with a garment reference, use **upper torso** as the target. "
                "For replacing an existing garment, use **shirt**, **white bikini**, **jacket**, etc."
            )

            result_video = gr.Video(
                label="Result"
            )

    # ========================================================
    # GENERATE
    # ========================================================


    generate_button.click(
        fn=process_video,
        inputs=[
            video_input,
            reference_input,
            target_input,
            prompt_input,
            strength_input,
            steps_input,
            guidance_input,
            seed_input,
        ],
        outputs=result_video,
        concurrency_limit=1,
    )


# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    # Download all large model artifacts BEFORE Gradio starts.
    # This means a user request can never trigger a fresh Wan download,
    # accepts work.
    ensure_model_downloads()

    print("All startup model downloads are complete.")
    print("Starting Gradio...")

    # IMPORTANT: HF Spaces may enable Gradio SSR by default.
    # For this long-running GPU job, SSR/Node proxying can misroute
    # heartbeat traffic as component payloads (e.g. "ping" -> FileData).
    # Disable SSR and use the normal client-rendered Gradio frontend.
    demo.queue(
        max_size=1,
    ).launch(
        css=CSS,
        ssr_mode=False,
    )