Spaces:
Running on Zero
Running on Zero
| """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) | |
| def conditioner(): | |
| return _remote(CONDITIONER_SPACE) | |
| 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 {}) | |
| 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")]) | |