# SPDX-License-Identifier: Apache-2.0 """Viggle Remake on ZeroGPU: re-render a video with its people (swap) or its scene (edit) changed. One click, three stages, and only the middle one holds the GPU: prepare (CPU) the source onto a 480-class canvas at 24 fps, its soundtrack cut out, the refs saved, and the prompt through MiniMax-H3's Qwen3-VL conditioner. H3 reads only `hidden_states[50]` of a text-only presentation, so `te_stream.py` runs embed_tokens and decoder layers 0-49 on the CPU, one layer at a time from /data-nvme, in a subprocess. No quota is spent on it; embeds are cached by token ids. render (GPU) 6 steps of the swap or edit LoRA on the stock MiniMax-H3 transformer, the refs pinned as keyframes at their frames, the source soundtrack pinned as a clean audio latent. finish (CPU) cut the frozen tail the 17n+5 grid forced, watermark, remux the source soundtrack. Host memory is the constraint, not the card: the container's cgroup holds 104 GB and the weights are ~83 GB. They are loaded with `disable_mmap=True` so they sit in anonymous memory. mmap-backed weights are page cache, and every encode streams ~54 GB of text-encoder shards through that same page cache -- it would evict them, and the next render's host-to-device copy would become a disk read. The load is at module level because `@spaces.GPU` bodies run in a worker forked from this process: it inherits what import time set up and nothing a request handler assigns later. """ import contextlib import hashlib import json import os import shutil import subprocess import sys import tempfile import time from pathlib import Path # 226 frames peaked at 91.0 GiB reserved on a B200 against this card's 94.4 GiB free; the 124-frame cap leaves # more room, and growable segments keep fragmentation out of it. Both spellings: torch 2.9 renamed the variable. os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import gradio as gr import numpy as np import spaces import torch from huggingface_hub import snapshot_download from PIL import Image H3 = "/tmp/h3" NVME = "/data-nvme" if os.path.isdir("/data-nvme") else "/tmp" # ~3 GB/s under memory pressure; overlay is 0.33 TE = f"{NVME}/h3/text_encoder" CACHE = Path(NVME) / "embeds" LORA = "/tmp/remake" HERE = Path(__file__).parent EX = HERE / "examples" MODEL_CARD = "https://huggingface.co/Viggle/Viggle-Remake" FPS = 24 CHUNK, TAIL = 17, 5 # 90 frames is the adapters' recommended minimum (the base pipeline's is 5 s). 124 (5.2 s) is the public cap: # GPU time per job. 226 fits in VRAM (87.2 GiB allocated / 91.0 reserved at 832x480), at about twice the time. MIN_FRAMES, MAX_FRAMES = 90, 124 MIN_UPLOAD_S = 2.0 # The 480-class canvas ladder the adapters run on, (w, h). # Every entry spends ~390 rows per latent frame, so the frame cap holds on all of them. LADDER = [(416, 960), (448, 896), (480, 832), (544, 736), (640, 640), (736, 544), (832, 480), (896, 448), (960, 416)] MAX_REFS = 3 SIGMAS = {"swap": [1, .990476, .965243, .923077, .8, .631579, 0], "edit": [1, .983607, .96, .923077, .857143, .705882, 0]} TEXT_TAG = 1 KEYFRAME_UNTIL = 0.9 # B200 timings for one adapter at 832x480, 6 steps including decode: 28.8 s at 124 frames, 59.5 s at 226. # SLOWDOWN is this card against that, H2D the weight copy onto it. Measured here 10-23 at 226 frames: 159-160 s # (2.69x) for both adapters, h2d 30-39 s; the defaults keep about 60 s of margin over that. SLOWDOWN = float(os.environ.get("SLOWDOWN", "3.0")) H2D = float(os.environ.get("H2D", "60")) ATTN = os.environ.get("ATTN", "_native_cudnn") # Animate measured it good on this card; "" for torch's default FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" assert Path(FONT).exists(), f"watermark font missing: {FONT}" t = time.time() snapshot_download("MiniMaxAI/MiniMax-H3", local_dir=H3, allow_patterns=["modular_model_index.json", "vae/*", "audio_vae/*", "scheduler/*", "audio_scheduler/*", "transformer/*", "processor/*", "tokenizer/*", "text_encoder/*.json"]) # Only the shards holding embed_tokens and layers 0-49: 53.7 of the text encoder's 66 GB. wmap = json.load(open(f"{H3}/text_encoder/model.safetensors.index.json"))["weight_map"] want = ("model.language_model.embed_tokens.", *[f"model.language_model.layers.{i}." for i in range(50)]) shards = sorted({f for k, f in wmap.items() if k.startswith(want)}) snapshot_download("MiniMaxAI/MiniMax-H3", local_dir=f"{NVME}/h3", allow_patterns=["text_encoder/*.json", *[f"text_encoder/{f}" for f in shards]]) snapshot_download("Viggle/Viggle-Remake", local_dir=LORA, allow_patterns=["lora/*"]) # public repo, no token print(f"[boot] weights in {time.time()-t:.0f}s ({len(shards)} text-encoder shards on {NVME})", flush=True) from diffusers import MiniMaxH3Transformer3DModel, ModularPipeline, attention_backend from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3ImageReference, MiniMaxH3VideoReference from diffusers.modular_pipelines.minimax_h3.before_denoise import (_ROPE_FRAME_RESCALE, MiniMaxH3Ref2VAPrepareLayoutStep, MiniMaxH3SetTimestepsStep, _frame_position_grid) from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3VideoDecodeStep from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3LoopSchedulerStep, MiniMaxH3Ref2VALoopDenoiser from diffusers.modular_pipelines.minimax_h3.encoders import (MiniMaxH3Ref2VAReferenceEncoderStep, MiniMaxH3Ref2VATextEncoderStep, encode_vae_condition) from diffusers.modular_pipelines.minimax_h3.modular_pipeline import audio_latent_num_frames from diffusers.utils.export_utils import encode_video from transformers import AutoProcessor, AutoTokenizer TOKENIZER = AutoTokenizer.from_pretrained(H3, subfolder="tokenizer") TPS = AutoProcessor.from_pretrained(H3, subfolder="processor").video_processor.temporal_patch_size # --------------------------------------------------------------------------------------------------------- # Per-request state the class patches below read at call time. A class patch is process-global and the GPU # worker serves every request, so nothing request-specific may be baked into a patch when it is installed. MODE = "swap" EMB = None # (1, T, 5120) bf16 prompt embeds KF = {"images": [], "anchors": []} # keyframe images on the canvas, and the output frame each is pinned at NKF = 0 # rows those keyframes take in the packed sequence def use_fixed_embeds(self, components, state): block_state = self.get_block_state(state) block_state.prompt_embeds = EMB.to(components._execution_device, torch.bfloat16) block_state.text_token_tags = torch.full((EMB.shape[1],), TEXT_TAG, dtype=torch.long) self.set_block_state(state, block_state) return components, state MiniMaxH3Ref2VATextEncoderStep.__call__ = use_fixed_embeds @torch.no_grad() def decode_to_uint8(self, components, state): """The stock decode step, but quantized on the card: the stock `postprocess_video` walks the clip to fp32 numpy on the CPU (~9 s of billed time at 243 frames) for frames that are bit-identical. From Viggle-Animate.""" block_state = self.get_block_state(state) device = components._execution_device latents_mean = torch.tensor(components.vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1) latents_std = torch.tensor(components.vae.config.latents_std, device=device).view(1, -1, 1, 1, 1) latents = block_state.latents * latents_std + latents_mean with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): video = components.vae.decode(latents, return_dict=False)[0] pixel_mean = torch.tensor(components.pixel_mean, device=device).view(1, -1, 1, 1, 1) pixel_std = torch.tensor(components.pixel_std, device=device).view(1, -1, 1, 1, 1) video = (video.float() * pixel_std + pixel_mean).clamp(0, 1) block_state.videos = [v.mul_(255).round_().to(torch.uint8).permute(1, 2, 3, 0).contiguous().cpu() for v in video] self.set_block_state(state, block_state) return components, state MiniMaxH3VideoDecodeStep.__call__ = decode_to_uint8 # Audio pin: the target audio rows hold the source soundtrack (or encoded silence) as a clean latent for the # whole denoise -- positional arg 6 is audio_timestep, and H3's t = 1 is clean -- and only video rows step. _build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(lambda *a: _build_row_timesteps(*a[:6], 1.0, *a[7:])) @torch.no_grad() def video_only_step(self, components, block_state, i, t): n = block_state.num_condition_video_rows block_state.latents[n:] = components.scheduler.step( block_state.noise_pred[0, n:].float(), t, block_state.latents[n:], return_dict=False)[0] return components, block_state MiniMaxH3LoopSchedulerStep.__call__ = video_only_step # Keyframes on a ref2va layout, as in the model card's inference/sample.py: each keyframe block is packed after # the references and pinned inside the target's rotary span at its real frame index. _encode_references = MiniMaxH3Ref2VAReferenceEncoderStep.__call__ def encode_references_then_keyframes(self, components, state): components, state = _encode_references(self, components, state) with torch.no_grad(): kf = [encode_vae_condition(components.vae, torch.from_numpy(np.array(im)).to(components._execution_device) .permute(2, 0, 1)[None, :, None], components.pixel_mean, components.pixel_std, components.keyframe_encode_seed) for im in KF["images"]] state.set("condition_latents", list(state.get("condition_latents") or []) + kf) return components, state MiniMaxH3Ref2VAReferenceEncoderStep.__call__ = encode_references_then_keyframes _layout = MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence def layout_with_keyframes(text_token_tags, references, condition_latents, audio_condition_latents, num_latent_frames, latent_height, latent_width, num_audio_latents, patch_size, audio_channels, audio_tag, video_tag): k = len(KF["anchors"]) pos, tags, vi, ai, ti, ncv, nca = _layout(text_token_tags, references, condition_latents[:len(condition_latents) - k], audio_condition_latents, num_latent_frames, latent_height, latent_width, num_audio_latents, patch_size, audio_channels, audio_tag, video_tag) if not k: return pos, tags, vi, ai, ti, ncv, nca s = text_token_tags.shape[0] + ncv + nca # the first target audio row; its time is the target's rotary origin origin = float(pos[s, 0]) blocks = [] for anchor, lat in zip(KF["anchors"], condition_latents[-k:]): grid, _ = _frame_position_grid(lat.shape[3], lat.shape[4], patch_size[1], patch_size[2]) p = torch.zeros(grid.shape[0], 3, dtype=torch.float64) p[:, 0] = origin + _ROPE_FRAME_RESCALE * float(anchor) p[:, 1:] = grid blocks.append(p) n = sum(b.shape[0] for b in blocks) pos = torch.cat([pos[:s], *blocks, pos[s:]]) tags = torch.cat([tags[:s], torch.full((n,), video_tag, dtype=tags.dtype), tags[s:]]) vi = torch.cat([vi[vi < s], torch.arange(s, s + n), vi[vi >= s] + n]) ai = torch.where(ai >= s, ai + n, ai) return pos, tags, vi, ai, ti, ncv + n, nca MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence = staticmethod(layout_with_keyframes) _denoise = MiniMaxH3Ref2VALoopDenoiser.__call__ def kf_release(self, components, block_state, i, t): """Once 1 - t reaches KEYFRAME_UNTIL the keyframe blocks are cut out of every per-row input.""" if not NKF or 1 - float(t) >= KEYFRAME_UNTIL: return _denoise(self, components, block_state, i, t) nkf, ncv = NKF, block_state.num_condition_video_rows f = block_state.denoiser_input_fields vi, ai = f["video_indices"], f["audio_indices"] lo, hi = int(vi[ncv - nkf]), int(vi[ncv - 1]) + 1 assert hi - lo == nkf and hi == int(ai[block_state.num_condition_audio_rows]), (lo, hi, nkf) keep = torch.ones(f["position_ids"].shape[0], dtype=torch.bool, device=vi.device) keep[lo:hi] = False plan = list(block_state.row_timestep_plan) plan[i] = (plan[i][0], plan[i][1][keep]) full = block_state.latents, f, block_state.row_timestep_plan block_state.latents = torch.cat([block_state.latents[:ncv - nkf], block_state.latents[ncv:]]) block_state.denoiser_input_fields = dict( f, position_ids=f["position_ids"][keep], token_tags=f["token_tags"][keep], video_indices=torch.cat([vi[:ncv - nkf], vi[ncv:] - nkf]), audio_indices=ai - nkf * (ai >= hi)) block_state.row_timestep_plan = plan components, block_state = _denoise(self, components, block_state, i, t) block_state.latents, block_state.denoiser_input_fields, block_state.row_timestep_plan = full p = block_state.noise_pred block_state.noise_pred = torch.cat([p[:, :ncv - nkf], p.new_zeros(1, nkf, p.shape[2]), p[:, ncv - nkf:]], 1) return components, block_state MiniMaxH3Ref2VALoopDenoiser.__call__ = kf_release # --------------------------------------------------------------------------------------------------------- t = time.time() PIPE = ModularPipeline.from_pretrained(H3, workflow="ref2va") type(PIPE).min_duration = property(lambda self: MIN_FRAMES / FPS) PIPE.load_components(names=["vae", "audio_vae", "scheduler", "audio_scheduler"], pretrained_model_name_or_path=H3, dtype=torch.bfloat16, disable_mmap=True) PIPE.transformer_ref = MiniMaxH3Transformer3DModel.from_pretrained(H3, subfolder="transformer", torch_dtype=torch.bfloat16, disable_mmap=True) # Both adapters stay loaded; `render` activates one and moves the other's 2.5 GiB off the card. # prefix=None: the saved keys are bare module names, and the default "transformer." filter loads nothing. for m in SIGMAS: PIPE.transformer_ref.load_lora_adapter(f"{LORA}/lora/{m}", weight_name="pytorch_lora_weights.safetensors", prefix=None, adapter_name=m) sh = lambda u, s: s * u / (1 + (s - 1) * u) # noqa: E731 GRIDS = {m: {s: [sh(x / (12 - 11 * x), s) for x in sig] for s in (12.0, 3.0)} for m, sig in SIGMAS.items()} for sched, s in ((PIPE.scheduler, 12.0), (PIPE.audio_scheduler, 3.0)): sched.set_timesteps = (lambda f, s: lambda *a, **k: f(sigmas=GRIDS[MODE][s], device=k.get("device")))( sched.set_timesteps, s) ASR = PIPE.audio_vae.config.sampling_rate def cgroup(): st = dict(line.split() for line in open("/sys/fs/cgroup/memory.stat")) if os.path.exists( "/sys/fs/cgroup/memory.stat") else {} return f"anon {int(st.get('anon', 0))/1e9:.1f} GB, file {int(st.get('file', 0))/1e9:.1f} GB" print(f"[boot] pipeline in {time.time()-t:.0f}s, {cgroup()}, disk free " f"{shutil.disk_usage('/tmp').free/1e9:.0f} GB /tmp, {shutil.disk_usage(NVME).free/1e9:.0f} GB {NVME}", flush=True) def run(*cmd): p = subprocess.run(cmd, capture_output=True, text=True, errors="replace") if p.returncode: raise gr.Error(f"{Path(cmd[0]).name} failed: " f"{p.stderr.strip().splitlines()[-1] if p.stderr.strip() else p.returncode}") return p.stdout def sweep(hours=6): """Drop working directories older than `hours`; age is the one property that cannot mistake a live job.""" cutoff = time.time() - hours * 3600 for d in Path(tempfile.gettempdir()).glob("tmp*"): with contextlib.suppress(OSError): if d.is_dir() and d.stat().st_mtime < cutoff: shutil.rmtree(d, ignore_errors=True) def probe(path): s = json.loads(run("ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height", "-of", "json", path))["streams"][0] return s["width"], s["height"] def nframes(path): return int(run("ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames", "-show_entries", "stream=nb_read_frames", "-of", "csv=p=0", path).strip()) def has_audio(path): return bool(run("ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries", "stream=codec_type", "-of", "csv=p=0", path).strip()) def grid(n): """Smallest 17k+5 frame count covering `n`; the tail is a frozen frame, cut off again in `finish`.""" return min(MAX_FRAMES, max(MIN_FRAMES, -(-(n - TAIL) // CHUNK) * CHUNK + TAIL)) def canvas(w, h): return min(LADDER, key=lambda b: abs(b[0] / b[1] - w / h)) def keyframe(path, w, h, first): """Onto the canvas the way our production renderer does: the first keyframe stretched, the rest cover-cropped.""" im = Image.open(path).convert("RGB") if im.size == (w, h): return im if first: return PIPE.image_processor.resize(im, height=h, width=w) scale = max(w / im.size[0], h / im.size[1]) size = (max(w, round(im.size[0] * scale)), max(h, round(im.size[1] * scale))) left, top = max(0, (size[0] - w) // 2), max(0, (size[1] - h) // 2) return im.resize(size, Image.Resampling.LANCZOS).crop((left, top, left + w, top + h)) def embed(frames, nref, prompt): """Prompt embeds for this presentation: `