"""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")
# Pika H3 Flex: a rank-128 LoRA on the same transformer_ref, sampled on its own fixed six-step re-noising schedule
# (from ref2va/manifest.json). Generic N-step Euler is not equivalent, so its step count is not a knob.
PIKA_REPO = os.environ.get("H3_PIKA_REPO", "Pika-Labs/Pika-H3-Flex")
PIKA_FILE = os.environ.get("H3_PIKA_FILE", "ref2va/adapter.safetensors")
PIKA_STEP_LIST = [1000, 750, 500, 250, 100, 40, 0]
PIKA_EVALS = len(PIKA_STEP_LIST) - 1
MODEL_PIKA, MODEL_TURBO = "Pika H3 Flex (6 steps)", "H3 Turbo LoRA"
DEFAULT_MODEL = MODEL_PIKA
DEFAULT_PIKA_HARMONIZE_STEP = 4  # of 6 evaluations: the subject free-refines over the last two (sigma 0.57 -> 0.33)
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"
DEFAULT_HARMONIZE_STEP = 6  # harmonize: drop the mask at this denoise step by default (of 8; later = more faithful, earlier = more relight)
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
PIKA_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, PIKA_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)
        try:
            # alpha-less file -> the loader pins alpha == rank (scale 1.0), Pika's own convention
            pipe.load_lora_weights(hf_hub_download(PIKA_REPO, PIKA_FILE), adapter_name="pika")
            PIKA_READY = True
            print(f"[bgswap] Pika H3 Flex ref2va adapter ready ({PIKA_REPO})", flush=True)
        except Exception as error:
            traceback.print_exc()
            print(f"[bgswap] Pika adapter 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. an example or an upload), a
    path relative to the app dir, and an external http(s) URL (downloaded here with the Space token)."""
    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.")

    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

    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")
    pika = request.get("model") == MODEL_PIKA
    adapter = "pika" if pika else "turbo"
    try:
        PIPE.enable_lora()
        PIPE.set_adapters([adapter], [1.0])
    except Exception as error:
        raise gr.Error(f"Could not switch to the {adapter} LoRA ({type(error).__name__}: {error})")

    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=(len(PIKA_STEP_LIST) if pika else int(steps)),
        generator=torch.Generator("cpu").manual_seed(int(seed)))
    if pika:
        kwargs["denoising_step_list"] = PIKA_STEP_LIST
        kwargs["stochastic_sampling"] = True
    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. The drop
        # step is user-picked and clamped below the step count (a drop at or past the last step cannot hand off).
        # `steps` is the model-evaluation count here; the handoff re-noises to timesteps[drop], so drop <= evals - 1.
        drop = int(request.get("harmonize_step") or DEFAULT_PIKA_HARMONIZE_STEP)
        kwargs["mask_drop_step"] = max(1, min(int(steps) - 1, drop))

    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,
                       harmonize_step=DEFAULT_PIKA_HARMONIZE_STEP, remove_only=False, progress=None, model=DEFAULT_MODEL):
    """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)
        model = model if model in (MODEL_PIKA, MODEL_TURBO) else DEFAULT_MODEL
        if model == MODEL_PIKA and not PIKA_READY:
            raise gr.Error("The Pika H3 Flex adapter failed to load; pick the H3 Turbo LoRA.")
        if model == MODEL_TURBO and not TURBO_READY:
            raise gr.Error("The turbo LoRA failed to load; pick Pika H3 Flex.")

    from h3_inpaint_crop import canvas_for_box, paste_back

    progress(0.1, desc="Removing the background (LTX 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.4, desc="Generating the new background (MiniMax-H3) ...")
    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, "model": model,
        "harmonize": bool(harmonize), "harmonize_step": int(harmonize_step),
    }
    # model evaluations: Pika's fixed six; turbo's `steps` sigma grid points drive steps - 1 of them
    evals = PIKA_EVALS if model == MODEL_PIKA else max(1, int(steps) - 1)
    progress(0.5, desc=f"Generating the new background ({model}, {width}x{height}, ~1-2 min) ...")
    generated, audio, sampling_rate = _generate(request, evals, 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)


# --------------------------------------------------------------------------------------------- UI (gr.Blocks)
_HERE = os.path.dirname(os.path.abspath(__file__))
EXAMPLES = [
    [os.path.join(_HERE, "examples/subject.mp4"),
     "a sunny tropical beach, white sand, palm trees and turquoise ocean waves behind the subject, bright warm daylight, static camera"],
    [os.path.join(_HERE, "examples/subject.mp4"),
     "a neon-lit cyberpunk city street at night, rain and glowing signs reflecting on wet asphalt behind the subject, cool moody light"],
    [os.path.join(_HERE, "examples/dance.mp4"),
     "a concert stage with dramatic blue and magenta spotlights and haze, dark background, cinematic"],
]

def run_classic(clip, prompt, budget, steps, seed, keep_audio, harmonize, harmonize_step, remove_only,
                model=DEFAULT_MODEL, progress=gr.Progress()):
    """Wraps replace_background so gradio injects a progress tracker - the handler then reports its stages
    (removing the background, then generating the new one) on the UI progress bar."""
    return replace_background(clip, prompt, budget, steps, seed, keep_audio, harmonize, harmonize_step,
                              remove_only, progress=progress, model=model)


with gr.Blocks(title="Video Replace Background") as demo:
    gr.Markdown(
        "# 🎬 Video Replace Background\n"
        "Drop your subject into a brand-new world. The "
        "[LTX-2.5 Alpha-Gen matte](https://huggingface.co/Lightricks/LTX-2.5-22b-IC-LoRA-Alpha-Gen) lifts them out, "
        "[Pika H3 Flex](https://huggingface.co/Pika-Labs/Pika-H3-Flex) on "
        "[MiniMax-H3](https://huggingface.co/multimodalart/MiniMax-H3-Pruned) "
        "([inpainting blocks](https://huggingface.co/diffusers-modular/minimax-h3-inpainting); the "
        "[turbo LoRA](https://huggingface.co/lightx2v/Minimax-h3-Turbo) is in Settings) paints the new scene "
        "around them in six steps - or just "
        "remove the background. **Harmonize** melts the subject's lighting into the scene so they don't look pasted on."
    )
    with gr.Row():
        with gr.Column():
            clip = gr.Video(label="Clip with a clear subject")
            remove_bg = gr.Checkbox(value=False, label="Just remove the background (skip generating a new one)")
            prompt = gr.Textbox(label="New background", placeholder="a sunny tropical beach, palm trees, ocean waves")
            harmonize = gr.Checkbox(value=True, label="Harmonize lighting into the scene",
                                    info="Blends the subject into the new background (recommended).")
            harmonize_step = gr.Slider(
                1, PIKA_EVALS - 1, value=DEFAULT_PIKA_HARMONIZE_STEP, step=1,
                label="Harmonize strength: drop mask at step",
                info="Out of the model's denoising steps. Lower = stronger relight, higher = truer to the original.")
            with gr.Accordion("Settings", open=False):
                model = gr.Radio([MODEL_PIKA, MODEL_TURBO], value=DEFAULT_MODEL, label="Model")
                budget = gr.Dropdown(list(BUDGETS), value=DEFAULT_BUDGET, label="Quality / speed")
                steps = gr.Slider(4, 28, value=DEFAULT_STEPS, step=1, label="Steps (turbo)",
                                  visible=DEFAULT_MODEL == MODEL_TURBO)
                seed = gr.Number(value=7, precision=0, label="Seed")
                keep_audio = gr.Checkbox(value=True, label="Keep the original audio")
            run = gr.Button("Go", variant="primary")
        with gr.Column():
            result = gr.Video(label="Result")
            with gr.Accordion("What the matte and mask looked like", open=False):
                matte_view = gr.Video(label="Alpha matte")
                mask_view = gr.Video(label="Background mask (white = repainted)")

    gr.Examples(examples=EXAMPLES, inputs=[clip, prompt], label="Examples", cache_examples=False)

    # the drop-step slider only matters while harmonize is on
    harmonize.change(lambda on: gr.update(visible=bool(on)), harmonize, harmonize_step)

    # the drop step must sit below the model-evaluation count (Pika: 6; turbo: Steps - 1), so the last allowed drop is
    # one before it: track that maximum and clamp the value to it
    def _drop_range(m, s, drop):
        top = PIKA_EVALS - 1 if m == MODEL_PIKA else max(1, int(s) - 2)
        return gr.update(maximum=top, value=min(int(drop), top))

    steps.change(_drop_range, [model, steps, harmonize_step], harmonize_step)
    model.change(_drop_range, [model, steps, harmonize_step], harmonize_step)
    model.change(lambda m: gr.update(visible=m == MODEL_TURBO), model, steps)

    run.click(run_classic, [clip, prompt, budget, steps, seed, keep_audio, harmonize, harmonize_step, remove_bg, model],
              [result, matte_view, mask_view], api_name="replace_background")


if __name__ == "__main__":
    load_pipeline()
    demo.launch(show_error=True, allowed_paths=[OUTPUT_DIR, os.path.join(_HERE, "examples")])