"""Put a video subject on a new, generated background.

Two models, chained. The subject is lifted with the LTX-2.5 Alpha-Gen matte LoRA (video in, alpha matte out),
reached over the gradio API on `hugging-apps/ltx-2-5-alpha-matte`. The background is then regenerated with MiniMax-H3
ref2va inpainting, held in this Space: the subject is handed to H3 as conditioning while the inverted matte marks the
background as the region to repaint. Text encoding runs remotely on the qwen3vl conditioner, the same split the
inpainting Space uses, so this Space never loads a 60 GiB encoder just to read a sentence.

Relight strength is the one real quality lever. At 0 the subject is held exactly (best when the new background's
light matches the original). Above 0 the subject's rows are given a small mask value so they partially denoise under
the new scene, which lets H3 warm, cool or rim-light the subject into a background whose light is different, while
keeping identity. At any relight above 0 the subject is read from H3's own output rather than pasted back, because
pasting the original pixels back is exactly what would undo the relight.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import time
import traceback
from functools import cache

import spaces  # patches torch.cuda so the model can load at startup rather than on GPU time
import gradio as gr
import numpy as np
from PIL import Image, ImageFilter

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "multimodalart/MiniMax-H3-Pruned")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
MATTE_SPACE = os.environ.get("H3_MATTE_SPACE", "hugging-apps/ltx-2-5-alpha-matte")
TURBO_REPO = os.environ.get("H3_TURBO_REPO", "lightx2v/Minimax-h3-Turbo")
TURBO_FILE = os.environ.get("H3_TURBO_FILE", "minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors")
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge").strip() or None
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
PLACEMENT = os.environ.get("H3_PLACEMENT", "eager").lower()  # eager: place on cuda at startup (ZeroGPU best practice)

FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
CANVAS_MULTIPLE = 32
MIN_DURATION = 2
DEFAULT_STEPS = 8
GROW_INTO_SUBJECT = 7  # odd MaxFilter kernel: eat ~3px of the matte edge halo back into the background

BUDGETS = {
    "Balanced · 640": (640, 640 * 1120),
    "Full · 768 short edge": (768, 768 * 1344),
    "Fast · 512": (512, 512 * 896),
}
DEFAULT_BUDGET = "Balanced · 640"
HARMONIZE_DROP_FRAC = 0.625  # harmonize: drop the mask after this fraction of the denoise steps (k5 of 8: most faithful + lighting)
OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "bgswap-outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Aspect-only labels the conditioner's canvas table knows; the request picks whichever is closest. The generation
# canvas itself comes from the budget, not from this label (ref2va reads only the prompt and the references).
CANVASES = {
    "960x544 · 16:9 fast": (544, 960), "1152x640 · 16:9": (640, 1152), "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344), "640x1152 · 9:16": (1152, 640), "768x1344 · 9:16 full": (1344, 768),
    "768x768 · 1:1 full": (768, 768), "1024x768 · 4:3 full": (768, 1024), "768x1024 · 3:4 full": (1024, 768),
    "1536x672 · 21:9 full": (672, 1536),
}

# Booking estimate, lifted from the inpainting Space's calibrated model (linear in rows for the matmuls, quadratic
# for attention, plus fixed encode/decode/placement terms).
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
CANVAS_SHORT_EDGE, CANVAS_MAX_PIXELS = 768, 768 * 1344
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
SOURCE_ENCODE_PER_DEFAULT_CANVAS = 12
MIN_GPU_DURATION, MAX_GPU_DURATION = 120, 1500

PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
TURBO_READY = False


# --------------------------------------------------------------------------------------------- model + remote halves
def load_pipeline() -> str | None:
    """Load the H3 ref2va inpaint generator + turbo at startup. The 60 GiB text encoder is never touched here."""
    global PIPE, MANAGER, LOAD_ERROR, TURBO_READY
    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR
    started = time.time()
    try:
        import torch
        from diffusers import ComponentsManager
        from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
        from huggingface_hub import hf_hub_download
        from h3_inpaint_blocks import MiniMaxH3Ref2VAInpaintGeneratorBlocks

        MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(MIN_DURATION))
        manager = ComponentsManager()
        pipe = MiniMaxH3Ref2VAInpaintGeneratorBlocks().init_pipeline(
            MODEL_REPO, components_manager=manager, collection="h3")
        pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer_ref.set_attention_backend(ATTENTION)
        try:
            pipe.load_lora_weights(hf_hub_download(TURBO_REPO, TURBO_FILE), adapter_name="turbo")
            pipe.set_adapters(["turbo"], [1.0])
            TURBO_READY = True
            print(f"[bgswap] turbo LoRA ready ({TURBO_REPO})", flush=True)
        except Exception as error:
            print(f"[bgswap] turbo unavailable ({type(error).__name__}: {error})", flush=True)
        # Place the whole pipeline (incl. the turbo LoRA) on cuda at startup. ZeroGPU's CUDA emulation handles
        # `.to("cuda")` outside the GPU function and startup placement is much faster than deferring it to the first
        # request - so the first generation is no longer penalised by a cold transfer. H3_PLACEMENT=lazy opts out.
        if PLACEMENT != "lazy":
            pipe.to("cuda")
            print("[bgswap] pipeline placed on cuda at startup", flush=True)
        PIPE, MANAGER = pipe, manager
        print(f"[bgswap] generator ready in {time.time() - started:.0f}s", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = f"**Loading `{MODEL_REPO}` failed**: `{type(error).__name__}: {error}`"
    return LOAD_ERROR


def _remote(space: str):
    from gradio_client import Client

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if os.environ.get("H3_REMOTE_TOKEN", "").strip() in ("1", "true", "yes") and token:
        return Client(space, token=token)
    return Client(space)


@cache
def conditioner():
    return _remote(CONDITIONER_SPACE)


@cache
def matte_client():
    return _remote(MATTE_SPACE)


def matte_remote(clip_path: str, seed: int):
    """`/generate_matte` on the Alpha-Gen Space: (subject-on-green, alpha matte, seed, details)."""
    from gradio_client import handle_file

    green, matte, seed_used, details = matte_client().predict(
        video_path=handle_file(clip_path), seed=float(seed), randomize_seed=False, api_name="/generate_matte")
    return green, matte, details


_MATTE_CACHE: dict[str, tuple[str, str]] = {}


def _clip_hash(path: str) -> str:
    import hashlib

    h = hashlib.md5()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def cached_matte(clip_path: str, seed: int):
    """The LTX matte pass is a function of the input video alone, so cache it by the clip's content. Changing only the
    prompt (or relight/steps/seed/background) then reuses the matte and skips the LTX pass and its ZeroGPU booking."""
    key = _clip_hash(clip_path)
    hit = _MATTE_CACHE.get(key)
    if hit and os.path.exists(hit[0]) and os.path.exists(hit[1]):
        print("[bgswap] reusing the cached matte for this clip (skipping the LTX pass)", flush=True)
        return hit
    green, matte, _ = matte_remote(clip_path, int(seed))
    # persist into the served output dir so the cache survives the gradio_client temp cleanup
    stamp = int(time.time() * 1000)
    green_out = os.path.join(OUTPUT_DIR, f"green-{stamp}.mp4")
    matte_out = os.path.join(OUTPUT_DIR, f"matte-src-{stamp}.mp4")
    shutil.copy(green, green_out)
    shutil.copy(matte, matte_out)
    _MATTE_CACHE[key] = (green_out, matte_out)
    return green_out, matte_out


def encode_remote(prompt, references, canvas, num_frames):
    """`/encode_ref2va` on the conditioner Space: prompt_embeds + text_token_tags in a safetensors file."""
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        media=[handle_file(p) for _, p in references],
        kinds=",".join(kind for kind, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=False,
        api_name="/encode_ref2va",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


# ------------------------------------------------------------------------------------------------------ pure helpers
def _resolve_clip(clip) -> str:
    """A clip input into a local filepath the matte Space can be handed.

    Handles every shape the UI can produce: a FileData dict (uploaded file), a plain path, a gradio file-serving
    URL (`/gradio_api/file=...`, which points at a file already on THIS Space - e.g. the pre-filled example or an
    upload), a path relative to the app dir, and an external http(s) URL (downloaded here with the Space token so a
    token-gated source works)."""
    import tempfile as _tf

    if isinstance(clip, dict):
        clip = clip.get("path") or clip.get("url")
    if not isinstance(clip, str) or not clip:
        raise gr.Error("Upload a clip with a clear subject.")

    # a gradio file URL -> the path it serves, which is on this Space's own disk
    if "/gradio_api/file=" in clip:
        clip = clip.split("/gradio_api/file=", 1)[1]

    if clip.startswith(("http://", "https://")):
        import requests

        tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        headers = {"Authorization": f"Bearer {tok}"} if tok else {}
        resp = requests.get(clip, headers=headers, timeout=120)
        resp.raise_for_status()
        local = os.path.join(_tf.mkdtemp(), "clip.mp4")
        with open(local, "wb") as handle:
            handle.write(resp.content)
        return local

    # a local path: resolve a relative one (e.g. "examples/subject.mp4") against the app dir
    if not os.path.isabs(clip) and not os.path.exists(clip):
        candidate = os.path.join(os.path.dirname(os.path.abspath(__file__)), clip)
        if os.path.exists(candidate):
            return candidate
    return clip


def read_video(path: str, limit: int | None = None) -> np.ndarray:
    import av

    with av.open(path) as container:
        frames = np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)])
    return frames if limit is None else frames[:limit]


def read_audio(path: str):
    import av
    import torch

    with av.open(path) as container:
        stream = next((s for s in container.streams if s.type == "audio"), None)
        if stream is None:
            return None, None
        chunks = [f.to_ndarray() for f in container.decode(audio=0)]
        if not chunks:
            return None, None
        waveform = torch.from_numpy(np.concatenate(chunks, axis=-1)).float()
        if waveform.shape[0] == 1 and stream.channels == 2:
            waveform = waveform.reshape(2, -1)
        return waveform, stream.rate


def floor_frames(available: int) -> int:
    """The largest `17 * n + 5 <= available` frame count the video VAE can decode."""
    n = (available - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK
    return max(LATENTS_PER_CHUNK, n * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK)


def background_mask(matte: np.ndarray, relight: float) -> np.ndarray:
    """From the white-is-subject matte: invert to a white-is-background repaint mask, harden, grow a few pixels into
    the subject to repaint the matte's edge halo. A relight above 0 gives the preserved subject interior that small
    value, so its rows partially denoise under the new scene instead of staying frozen."""
    subject = (matte.mean(-1) / 255.0 > 0.5).astype(np.float32)
    bg = 1.0 - subject
    grown = np.stack([
        np.asarray(Image.fromarray((f * 255).astype(np.uint8)).filter(ImageFilter.MaxFilter(GROW_INTO_SUBJECT)),
                   dtype=np.float32) / 255.0
        for f in bg
    ])
    bg = (grown > 0.5).astype(np.float32)
    if relight > 0:
        bg = np.where(bg > 0.5, 1.0, float(relight)).astype(np.float32)
    return bg


def gray_reference(green0: np.ndarray, matte0: np.ndarray) -> str:
    """The subject over neutral gray, from frame 0: a clean identity reference with no background to bias the new one."""
    alpha = (matte0.mean(-1) / 255.0 > 0.5).astype(np.float32)[..., None]
    ref = (green0 * alpha + 128 * (1 - alpha)).round().clip(0, 255).astype(np.uint8)
    path = os.path.join(tempfile.mkdtemp(), "subject-reference.png")
    Image.fromarray(ref).save(path)
    return path


def nearest_canvas_label(height: int, width: int) -> str:
    target = width / height
    return min(CANVASES, key=lambda label: abs(CANVASES[label][1] / CANVASES[label][0] - target))


def mask_to_clip(mask: np.ndarray, stamp: int) -> str:
    """Write a grayscale mask preview clip (for the UI), into the served output dir."""
    from diffusers.utils import encode_video
    import torch

    rgb = np.stack([np.stack([(f * 255).astype(np.uint8)] * 3, -1) for f in mask])
    path = os.path.join(OUTPUT_DIR, f"mask-{stamp}.mp4")
    encode_video(torch.from_numpy(np.ascontiguousarray(rgb)), fps=FPS, output_path=path)
    return path


def video_latent_frames(num_frames: int) -> int:
    return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2


def target_rows(height: int, width: int, num_frames: int) -> int:
    video = video_latent_frames(num_frames) * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS


def reference_rows(references, num_frames, image_edge=768) -> int:
    rows = 0
    for kind, path in references:
        if kind == "image":
            width, height = Image.open(path).size
            scale = min(1.0, image_edge / min(width, height))
            resolved = [max(CANVAS_MULTIPLE, round(e * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
                        for e in (height, width)]
            rows += (resolved[0] // CANVAS_MULTIPLE) * (resolved[1] // CANVAS_MULTIPLE)
    return rows


def get_duration(request, steps, seed, **_):
    height, width, num_frames = request["height"], request["width"], request["num_frames"]
    references = request["references"]
    sequence = (int(request["text_token_tags"].shape[0]) + reference_rows(references, num_frames)
                + target_rows(height, width, num_frames))
    denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence ** 2) * SAFETY
    encode = 5 + reference_rows(references, num_frames) * 1e-3
    share = (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
    # Only budget the cuda-transfer allowance when placement is deferred to the GPU call; eager placement pays it
    # once at startup, off the clock, so the per-request reservation (and queue priority) improves.
    placement = PLACEMENT_ALLOWANCE if PLACEMENT == "lazy" else 0
    total = (placement + encode + SOURCE_ENCODE_PER_DEFAULT_CANVAS * share + denoise
             + DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * share + 10)
    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    print(f"[bgswap] S={sequence} -> reserving {duration}s ({denoise:.0f}s denoise at {steps} steps)", flush=True)
    return duration


def build_references(references, image_edge: int = 768):
    from diffusers.modular_pipelines.minimax_h3.references import MiniMaxH3ImageReference

    built = []
    for kind, path in references:
        image = Image.open(path).convert("RGB")
        scale = image_edge / min(image.size)
        if scale < 1.0:
            image = image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS)
        built.append(MiniMaxH3ImageReference(image=np.asarray(image)))
    return built


def check_prompt(prompt: str) -> None:
    """Prompt-level safety gate, soft: if the guard module or its backend is unavailable, allow and log."""
    try:
        import ncii_guard

        ncii_guard.check_prompt(prompt)
    except Exception as error:
        print(f"[bgswap] prompt guard skipped ({type(error).__name__}: {error})", flush=True)


# ------------------------------------------------------------------------------------------------------- GPU + handler
GPU_KWARGS = {"duration": get_duration} | ({"size": GPU_SIZE} if GPU_SIZE else {})


@spaces.GPU(**GPU_KWARGS)
def _generate(request, steps, seed):
    import torch

    if PLACEMENT == "lazy":
        PIPE.to("cuda")
    try:
        if request.get("turbo") and TURBO_READY:
            PIPE.enable_lora()
            PIPE.set_adapters(["turbo"], [1.0])
        elif TURBO_READY:
            PIPE.disable_lora()
    except Exception as error:
        print(f"[bgswap] turbo switch failed ({type(error).__name__}: {error})", flush=True)

    kwargs = dict(
        prompt_embeds=request["prompt_embeds"].to("cuda"), text_token_tags=request["text_token_tags"],
        references=build_references(request["references"]),
        source_video=request["source_video"], source_fps=FPS, mask=request["mask"],
        source_audio=request["source_audio"], source_audio_sample_rate=request["source_audio_sample_rate"],
        audio_mask=None, height=request["height"], width=request["width"], num_frames=request["num_frames"],
        num_inference_steps=int(steps), generator=torch.Generator("cpu").manual_seed(int(seed)))
    if request.get("harmonize"):
        # single pass: keep the mask until mask_drop_step, then free-refine the whole frame so the subject blends
        # into the scene (the subject is re-noised to that step's level at the handoff). No second pass.
        kwargs["mask_drop_step"] = max(1, min(int(steps) - 1, round(int(steps) * HARMONIZE_DROP_FRAC)))

    state = PIPE(**kwargs)
    audio = state.get("audio")
    return state.get("videos")[0], None if audio is None else audio[0].cpu(), state.get("sampling_rate")


# The matte Space composites with this chroma: keyed = alpha * subject + (1 - alpha) * GREEN_SCREEN (its own constant).
GREEN_SCREEN = np.array([0, 177, 64], dtype=np.float32)


def dekey(green: np.ndarray, matte: np.ndarray) -> np.ndarray:
    """Recover the subject on black from the matte Space's green composite - exactly, with no green anywhere.

    The Space builds `keyed = alpha * subject + (1 - alpha) * GREEN_SCREEN`, so subtracting the green term returns
    `alpha * subject` (the subject premultiplied on black). This is not a despill approximation: it removes the green
    contribution in closed form, so the subject the inpaint preserves carries true edges and no green rim can flicker.
    """
    n = min(len(green), len(matte))
    g = green[:n]
    # GREEN_SCREEN red is 0, so the R channel is unchanged; only G and B get the green term removed. The matte is
    # grayscale, so one channel is the alpha. Working per-channel avoids a full (n, H, W, 3) float broadcast.
    spill = 1.0 - matte[:n, ..., 0].astype(np.float32) / 255.0
    out = g.copy()
    out[..., 1] = np.clip(g[..., 1].astype(np.float32) - spill * GREEN_SCREEN[1], 0, 255).astype(np.uint8)
    out[..., 2] = np.clip(g[..., 2].astype(np.float32) - spill * GREEN_SCREEN[2], 0, 255).astype(np.uint8)
    return out


def replace_background(clip, prompt, budget, steps, seed, keep_audio, harmonize=True, remove_only=False, progress=None):
    """Clip + a background description -> the subject on that new background. With remove_only, the subject on black
    (no new background, no H3 - just the matte)."""
    progress = progress or (lambda *a, **k: None)
    remove_only = bool(remove_only)
    clip = _resolve_clip(clip)
    if not remove_only:
        if LOAD_ERROR:
            raise gr.Error(LOAD_ERROR)
        if PIPE is None:
            raise gr.Error("The generator is still loading.")
        if not (prompt or "").strip():
            raise gr.Error("Describe the background you want, or turn on 'Just remove the background'.")
        check_prompt(prompt)

    from h3_inpaint_crop import canvas_for_box, paste_back

    progress(0.05, desc="Lifting the subject (alpha matte) ...")
    green_path, matte_path = cached_matte(clip, int(seed))
    # Un-composite the green up front: `green` becomes the subject on black (green removed in closed form from the
    # matte Space's keyed output), so the subject the inpaint preserves is green-free and no rim can flicker.
    matte = read_video(matte_path)
    green = dekey(read_video(green_path), matte)

    if remove_only:
        import torch
        from diffusers.utils import encode_video

        # `green` is already the subject on black (dekeyed from the green composite), which is exactly what the
        # remove-background output is - no background generated, no H3 denoise.
        progress(0.6, desc="Removing the background ...")
        cutout = green
        waveform, sample_rate = (read_audio(clip) if keep_audio else (None, None))
        stamp = int(time.time() * 1000)
        out_path = os.path.join(OUTPUT_DIR, f"removed-{stamp}.mp4")
        encode_video(torch.from_numpy(np.ascontiguousarray(cutout)), fps=FPS, output_path=out_path,
                     audio=None if waveform is None else waveform, audio_sample_rate=sample_rate)
        matte_out = os.path.join(OUTPUT_DIR, f"matte-{stamp}.mp4")
        shutil.copy(matte_path, matte_out)
        return out_path, matte_out, matte_out
    n = floor_frames(min(len(green), len(matte)))
    green, matte = green[:n], matte[:n]
    clip_h, clip_w = green.shape[1], green.shape[2]

    mask = background_mask(matte, 0.0)[:n]
    ref_path = gray_reference(green[0], matte[0])

    short_edge, max_pixels = BUDGETS[budget]
    box = (0, 0, (clip_h // CANVAS_MULTIPLE) * CANVAS_MULTIPLE, (clip_w // CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
    height, width = canvas_for_box(box[2], box[3], CANVAS_MULTIPLE, short_edge, max_pixels)

    references = [("image", ref_path)]
    full_prompt = "<Picture 1> " + prompt.strip()
    progress(0.15, desc="Reading the prompt and reference ...")
    prompt_embeds, text_token_tags, metadata, _ = encode_remote(
        full_prompt, references, nearest_canvas_label(box[2], box[3]), n)
    num_frames = int(metadata["num_frames"])
    green, mask = green[:num_frames], mask[:num_frames]

    waveform, sample_rate = (read_audio(clip) if keep_audio else (None, None))
    request = {
        "prompt_embeds": prompt_embeds, "text_token_tags": text_token_tags, "references": references,
        "source_video": green, "mask": mask, "source_audio": waveform, "source_audio_sample_rate": sample_rate,
        "height": height, "width": width, "num_frames": num_frames, "turbo": True,
        "harmonize": bool(harmonize),
    }
    progress(0.25, desc=f"Repainting the background at {width}x{height} ...")
    generated, audio, sampling_rate = _generate(request, steps, seed)

    progress(0.95, desc="Finishing ...")
    raw = np.stack([np.asarray(f) for f in generated])
    gen = np.stack([np.asarray(Image.fromarray(f).resize((clip_w, clip_h), Image.Resampling.LANCZOS)) for f in raw])
    if bool(harmonize):
        # the harmonize (full-frame mask-drop) pass produced the final subject in H3's own frame; use it directly
        frames = gen
    else:
        # Full-frame composite: restore the exact subject pixels, feathering ONLY the subject boundary. paste_back's
        # border ramp fades to the plate at the frame edges, which here is the green screen - so it is not used; the
        # background (incl. the frame edge) is taken fully from the generated frame.
        from h3_inpaint_crop import _feather_mask

        weight = _feather_mask((mask > 0.5).astype(np.float32), 8)[..., None]
        frames = (green.astype(np.float32) * (1 - weight) + gen.astype(np.float32) * weight)
        frames = frames.round().clip(0, 255).astype(np.uint8)

    from diffusers.utils import encode_video
    import torch

    stamp = int(time.time() * 1000)
    out_path = os.path.join(OUTPUT_DIR, f"bgswap-{stamp}.mp4")
    encode_video(torch.from_numpy(np.ascontiguousarray(frames)), fps=FPS, output_path=out_path,
                 audio=None if audio is None else audio, audio_sample_rate=sampling_rate)
    matte_out = os.path.join(OUTPUT_DIR, f"matte-{stamp}.mp4")
    shutil.copy(matte_path, matte_out)
    return out_path, matte_out, mask_to_clip(mask, stamp)


# ----------------------------------------------------------------------------------------- flow UI (gr.Workflow)
def _filedict(path: str) -> dict:
    """A video file as the FileData dict a workflow subject renders."""
    return {"path": path, "url": f"/gradio_api/file={path}", "orig_name": os.path.basename(path),
            "mime_type": "video/mp4", "meta": {"_type": "gradio.FileData"}}


def generate_video(clip, prompt, budget=DEFAULT_BUDGET, steps=DEFAULT_STEPS, seed=7, keep_audio=True,
                   harmonize=True, remove_only=False):
    """The workflow's fn operator: a clip + a background description -> result, alpha matte, background mask. With
    harmonize on, a second mask-drop pass blends the subject's lighting into the scene. With remove_only on, just
    the subject on black (no new background, no H3)."""
    remove_only = bool(remove_only)
    if not remove_only:
        if LOAD_ERROR:
            raise gr.Error(LOAD_ERROR.replace("**", "").replace("`", ""))
        if PIPE is None:
            raise gr.Error("The generator is still loading - watch the Space logs and retry shortly.")
    try:
        result, matte, mask = replace_background(
            clip, prompt, str(budget or DEFAULT_BUDGET),
            int(steps or DEFAULT_STEPS), float(seed or 0), bool(keep_audio), bool(harmonize), remove_only)
    except gr.Error:
        raise
    except Exception as error:
        message = str(error).lower()
        if any(h in message for h in ("gpu limit", "quota", "could not allocate", "too many", "concurrent")):
            raise gr.Error("The shared ZeroGPU pool is at capacity right now. Wait a minute and retry.") from error
        raise
    return _filedict(result), _filedict(matte), _filedict(mask)


_HERE = os.path.dirname(os.path.abspath(__file__))
# gr.Workflow needs gradio >= 6.25; None under older local gradio so the self-test can still import this module.
demo = (gr.Workflow(graph=os.path.join(_HERE, "workflow.json"), bind={"generate_video": generate_video})
        if hasattr(gr, "Workflow") else None)


# ----------------------------------------------------------------------------------------------------------- self-test
if __name__ == "__main__":
    if demo is None:
        raise SystemExit("gr.Workflow requires gradio >= 6.25")
    load_pipeline()
    demo.launch(show_error=True, allowed_paths=[OUTPUT_DIR, os.path.join(_HERE, "examples")])