Spaces:
Running on Zero
Running on Zero
| # 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 | |
| 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:])) | |
| 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: `<Video 1>: `, a time label per 2 fps block of the source, one label | |
| per ref, then the prompt, each piece tokenized on its own. Cached by token ids; the examples ship theirs.""" | |
| _, stamps = MiniMaxH3Ref2VATextEncoderStep._sample_video_condition_frames( | |
| np.zeros((frames, 1, 1, 3), dtype=np.uint8), 24.0, 2.0, TPS) | |
| segments = ["<Video 1>: ", *[f"<{s:.1f} seconds>" for s in stamps], *[f"<Picture {i + 1}>: " for i in range(nref)], | |
| prompt] | |
| ids = [i for s in segments for i in TOKENIZER(s, add_special_tokens=False)["input_ids"]] | |
| name = hashlib.sha1(json.dumps(ids).encode()).hexdigest()[:20] + ".pt" | |
| if (EX / "embeds" / name).exists(): | |
| return str(EX / "embeds" / name), "".join(segments), 0.0 | |
| out = CACHE / name | |
| t0 = time.time() | |
| if not out.exists(): | |
| CACHE.mkdir(exist_ok=True) | |
| (CACHE / f"{name}.ids").write_text(json.dumps(ids)) | |
| run(sys.executable, str(HERE / "te_stream.py"), TE, str(CACHE / f"{name}.ids"), f"{out}.part") | |
| os.replace(f"{out}.part", out) | |
| print(f"[te] {len(ids)} tokens in {time.time()-t0:.0f}s, {cgroup()}", flush=True) | |
| return str(out), "".join(segments), time.time() - t0 | |
| def prepare(mode, video, prompt, refs, pins): | |
| if not video or not os.path.exists(video): | |
| raise gr.Error("upload a source video (or reload the page if this one predates a Space restart)") | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("write the prompt -- the examples show the six-section format it has to follow") | |
| if not refs: | |
| raise gr.Error("give at least one reference image") | |
| sweep() | |
| work = Path(tempfile.mkdtemp()) | |
| w, h = canvas(*probe(video)) | |
| full = str(work / "full.mp4") | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", video, "-vf", | |
| f"scale={w}:{h}:force_original_aspect_ratio=increase,crop={w}:{h},fps={FPS}", | |
| "-frames:v", str(MAX_FRAMES), "-c:v", "libx264", "-qp", "0", "-pix_fmt", "yuv420p", "-an", full) # lossless: one lossy encode, as sample.py | |
| have = nframes(full) | |
| if have < MIN_UPLOAD_S * FPS: | |
| raise gr.Error(f"the source is {have/FPS:.1f}s at {FPS} fps; at least {MIN_UPLOAD_S:.0f}s is needed") | |
| if any(not 0 <= p < have for p in pins): | |
| raise gr.Error(f"pin frames must be in 0..{have - 1} (the source has {have} frames at {FPS} fps)") | |
| frames, keep = grid(have), have / FPS | |
| cond = str(work / "cond.mp4") | |
| pad = ["-vf", "tpad=stop_mode=clone:stop=-1"] if have < frames else [] | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", full, *pad, "-frames:v", str(frames), "-c:v", "libx264", "-crf", "12", | |
| "-pix_fmt", "yuv420p", "-an", cond) | |
| apath = None | |
| if has_audio(video): | |
| apath = str(work / "audio.f32") | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", video, "-t", f"{keep:.3f}", "-vn", "-f", "f32le", "-ac", "2", | |
| "-ar", str(ASR), apath) | |
| # the refs are copied out of gradio's upload dir, which a restart empties; keyframes go onto the canvas | |
| saved, kfs = [], [] | |
| for i, r in enumerate(refs): | |
| saved.append(str(work / f"ref{i + 1}.png")) | |
| Image.open(r).convert("RGB").save(saved[-1]) | |
| if pins: | |
| kfs.append(str(work / f"kf{i + 1}.png")) | |
| keyframe(r, w, h, i == 0).save(kfs[-1]) | |
| emb, presentation, te_s = embed(frames, len(refs), prompt) | |
| return dict(mode=mode, cond=cond, w=w, h=h, frames=frames, keep=keep, apath=apath, refs=saved, kfs=kfs, | |
| pins=pins, emb=emb, presentation=presentation, te_s=te_s) | |
| def budget(job, seed): | |
| """Seconds asked of ZeroGPU: an admission check against the caller's quota (xlarge counts double) and a | |
| queue hint. Billing is on the actual window, so this only has to be a safe over-estimate.""" | |
| return round(30 + H2D + SLOWDOWN * (18.6 + 0.301 * (job["frames"] - MIN_FRAMES))) | |
| def render(job, seed): | |
| global MODE, EMB, NKF | |
| t0 = time.time() | |
| PIPE.to("cuda") | |
| other = "edit" if job["mode"] == "swap" else "swap" | |
| for m in PIPE.transformer_ref.modules(): | |
| if hasattr(m, "lora_A") and other in m.lora_A: | |
| m.lora_A[other].to("cpu") | |
| m.lora_B[other].to("cpu") | |
| torch.cuda.empty_cache() | |
| PIPE.transformer_ref.set_adapters(job["mode"]) | |
| h2d = time.time() - t0 | |
| torch.cuda.reset_peak_memory_stats() | |
| if ATTN: | |
| PIPE.transformer_ref.set_attention_backend(ATTN) # process-wide: the audio encode below opts out | |
| MODE, w, h, frames = job["mode"], job["w"], job["h"], job["frames"] | |
| EMB = torch.load(job["emb"]) | |
| KF["images"] = [Image.open(p).convert("RGB") for p in job["kfs"]] | |
| KF["anchors"] = list(job["pins"]) | |
| r, (_, ph, pw) = PIPE.vae_spatial_compression_ratio, PIPE.patch_size | |
| NKF = len(job["pins"]) * (h // r // ph) * (w // r // pw) | |
| PIPE.register_to_config(canvas_short_edge=min(w, h), canvas_max_pixels=w * h, reference_image_short_edge=min(w, h)) | |
| # the soundtrack (or silence) as clean target audio latents; cuDNN has no causal kernel for the audio encoder | |
| n_samp = round((frames + 1) / FPS * ASR) | |
| wav = (torch.frombuffer(bytearray(Path(job["apath"]).read_bytes()), dtype=torch.float32).view(-1, 2).T | |
| if job["apath"] else torch.zeros(2, 0)) | |
| wav = torch.nn.functional.pad(wav, (0, max(0, n_samp - wav.shape[1])))[:, :n_samp] | |
| with attention_backend(), torch.no_grad(): | |
| posterior = PIPE.audio_vae.encode(wav[:, None].to("cuda", torch.bfloat16), return_dict=False)[0] | |
| mean = torch.tensor(PIPE.audio_vae.config.latents_mean).view(1, 1, -1) | |
| std = torch.tensor(PIPE.audio_vae.config.latents_std).view(1, 1, -1) | |
| rows = (posterior.mode().float().cpu().transpose(1, 2)[:, :audio_latent_num_frames(frames)] - mean) / std | |
| video = MiniMaxH3VideoReference.from_file(job["cond"]) | |
| video.audio, video.sample_rate = None, None # the soundtrack is the target audio pin, not reference audio | |
| t1 = time.time() | |
| result = PIPE(prompt=job["presentation"], | |
| references=[video, *[MiniMaxH3ImageReference.from_file(x) for x in job["refs"]]], | |
| num_frames=frames, height=h, width=w, num_inference_steps=len(SIGMAS[MODE]) - 1, | |
| generator=torch.Generator().manual_seed(int(seed or 0)), | |
| audio_latents=rows.permute(0, 2, 1).contiguous(), output=["videos", "audio", "sampling_rate"]) | |
| took = time.time() - t1 | |
| raw = str(Path(job["cond"]).parent / "raw.mp4") | |
| encode_video(result["videos"][0], fps=FPS, output_path=raw, audio=result["audio"][0], | |
| audio_sample_rate=result["sampling_rate"]) | |
| msg = (f"{MODE}: {took:.0f}s for {len(SIGMAS[MODE]) - 1} steps at {w}x{h}x{frames}, peak " | |
| f"{torch.cuda.max_memory_allocated()/2**30:.1f} GiB allocated / " | |
| f"{torch.cuda.max_memory_reserved()/2**30:.1f} reserved, weights onto the card in {h2d:.0f}s") | |
| print(f"[gpu] {msg} (budget {budget(job, seed)}s)", flush=True) | |
| return raw, msg | |
| def finish(raw, job): | |
| """Cut the frozen tail, watermark, and remux the source soundtrack -- exact, where the render's own audio rows | |
| only come back through the audio VAE. A silent source comes back silent.""" | |
| out = str(Path(raw).parent / "viggle-remake.mp4") | |
| src = ["-f", "f32le", "-ar", str(ASR), "-ac", "2", "-i", job["apath"]] if job["apath"] else [] | |
| amap = ["-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac"] if job["apath"] else ["-an"] | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", raw, *src, "-t", f"{job['keep']:.3f}", "-vf", | |
| f"drawtext=fontfile={FONT}:text='viggle-remake':x=w-tw-14:y=h-th-14:fontsize=22:" | |
| "fontcolor=white@0.85:shadowcolor=black@0.6:shadowx=1:shadowy=1", | |
| *amap, "-crf", "18", "-pix_fmt", "yuv420p", out) | |
| return out | |
| def remake(mode, video, prompt, refs, pins, seed): | |
| """Prepare and finish run here, in the app's process; only `render` takes the GPU. `refs` is the gallery | |
| ((path, caption) pairs, in upload order), `pins` a comma list with one frame per picture, or empty.""" | |
| refs = [r[0] if isinstance(r, (list, tuple)) else r for r in refs or []] | |
| if len(refs) > MAX_REFS: | |
| raise gr.Error(f"at most {MAX_REFS} reference pictures") | |
| try: | |
| pins = [int(p) for p in (pins or "").replace(",", " ").split()] | |
| except ValueError: | |
| raise gr.Error("pin frames are whole numbers, comma separated, e.g. 12, 96") | |
| if pins and len(pins) != len(refs): | |
| raise gr.Error(f"{len(refs)} picture(s) but {len(pins)} pin frame(s): give one per picture, or none") | |
| t0 = time.time() | |
| job = prepare(mode, video, prompt, refs, pins) | |
| t1 = time.time() | |
| raw, msg = render(job, seed) | |
| out = finish(raw, job) | |
| te = f"prompt encoded in {job['te_s']:.0f}s" if job["te_s"] else "prompt embeds cached" | |
| return out, f"{msg}\nprepare {t1 - t0:.0f}s ({te}), whole job {time.time() - t0:.0f}s" | |
| CASES = { | |
| "swap-pineapple-mascot": ("swap", [12, 96], "the pizza maker becomes a plush pineapple mascot"), | |
| "edit-remove-turbines": ("edit", [12, 96], "every wind turbine removed from a drone shot"), | |
| } | |
| # Viggle look, as in the Viggle Image Swap Space: ink on cream, Inter, the "fine-tuned by Viggle" badge | |
| INK, CREAM = "#29231e", "#f4f1ec" | |
| THEME = gr.themes.Base( | |
| primary_hue=gr.themes.colors.stone, secondary_hue=gr.themes.colors.stone, neutral_hue=gr.themes.colors.stone, | |
| font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], | |
| radius_size=gr.themes.sizes.radius_md, | |
| ).set( | |
| body_background_fill=CREAM, body_background_fill_dark="#1c1814", | |
| body_text_color=INK, body_text_color_dark=CREAM, | |
| body_text_color_subdued="#6f655b", body_text_color_subdued_dark="#b3a99d", | |
| background_fill_primary="#fbf9f6", background_fill_primary_dark="#24201b", | |
| background_fill_secondary="#efe9e1", background_fill_secondary_dark="#2c2721", | |
| block_background_fill="#fbf9f6", block_background_fill_dark="#24201b", | |
| block_border_color="#e2dbd1", block_border_color_dark="#3a332c", | |
| border_color_primary="#e2dbd1", border_color_primary_dark="#3a332c", | |
| block_label_text_color="#6f655b", block_label_text_color_dark="#b3a99d", | |
| block_title_text_color=INK, block_title_text_color_dark=CREAM, | |
| block_shadow="none", | |
| input_background_fill="#ffffff", input_background_fill_dark="#1c1814", | |
| input_border_color="#e2dbd1", input_border_color_dark="#3a332c", | |
| input_border_color_focus="#8a7b6c", input_border_color_focus_dark="#a8988a", | |
| button_primary_background_fill=INK, button_primary_background_fill_hover="#3f362e", | |
| button_primary_text_color=CREAM, button_primary_border_color=INK, | |
| button_primary_background_fill_dark=CREAM, button_primary_background_fill_hover_dark="#e3dbcf", | |
| button_primary_text_color_dark=INK, button_primary_border_color_dark=CREAM, | |
| button_secondary_background_fill="#efe9e1", button_secondary_background_fill_dark="#2c2721", | |
| button_secondary_text_color=INK, button_secondary_text_color_dark=CREAM, | |
| color_accent=INK, color_accent_soft="#e9e2d8", color_accent_soft_dark="#3a332c", | |
| slider_color=INK, slider_color_dark=CREAM, | |
| checkbox_background_color_selected=INK, checkbox_background_color_selected_dark="#a08a72", | |
| checkbox_border_color_dark="#5a5047", | |
| checkbox_border_color_selected=INK, checkbox_border_color_selected_dark="#a08a72", | |
| checkbox_label_background_fill_selected="#e9e2d8", checkbox_label_background_fill_selected_dark="#3a332c", | |
| checkbox_label_text_color_selected=INK, checkbox_label_text_color_selected_dark=CREAM, | |
| link_text_color=INK, link_text_color_hover="#6f655b", link_text_color_visited=INK, | |
| link_text_color_dark=CREAM, link_text_color_hover_dark="#b3a99d", link_text_color_visited_dark=CREAM, | |
| ) | |
| # The reference gallery (elem_id "refs"), as in the Viggle Qwen-Image turbo Space. Gradio's gallery takes drops only | |
| # while it is empty; REFS_HEAD forwards a drop on the filled gallery to its upload input (which appends), and | |
| # "+ Add picture" clicks it. REFS_CSS keeps three columns at any count, numbers the thumbnails "<Picture N>" in upload | |
| # order (the order the prompt and the pin list refer to), keeps the delete button visible, turns off thumbnail clicks | |
| # (their preview hides the upload icon), and shows the add button only while the gallery holds 1..MAX_REFS-1 pictures. | |
| REFS_HEAD = """<script> | |
| document.addEventListener("dragover", (e) => { if (e.target.closest?.("#refs .grid-wrap")) e.preventDefault(); }, true); | |
| document.addEventListener("drop", (e) => { | |
| const input = e.target.closest?.("#refs")?.querySelector(".grid-wrap input[type=file]"); | |
| if (!input || !e.dataTransfer?.files.length) return; | |
| e.preventDefault(); | |
| e.stopPropagation(); | |
| input.files = e.dataTransfer.files; | |
| input.dispatchEvent(new Event("change")); | |
| }, true); | |
| </script>""" | |
| ADD_REF_JS = "() => document.querySelector('#refs input[type=file]')?.click()" | |
| CSS = """ | |
| .gradio-container {max-width: 1320px !important; margin: 0 auto !important;} | |
| .dark .gradio-container {--color-accent: #f4f1ec;} | |
| #hdr {display: flex; align-items: center; gap: 18px; flex-wrap: wrap; margin: 6px 0 2px;} | |
| #hdr h1 {font-size: 30px; font-weight: 600; letter-spacing: -0.015em; margin: 0; color: var(--body-text-color);} | |
| #hdr .badge {display: inline-flex; color: var(--body-text-color);} | |
| #hdr .badge svg {height: 30px; width: auto;} | |
| #intro p {max-width: 80ch; line-height: 1.55;} | |
| #refs .grid-wrap {min-height: 0; max-height: none; padding-top: 34px;} | |
| #refs .grid-container {grid-template-columns: repeat(3, minmax(0, 1fr)); counter-reset: ref;} | |
| #refs .gallery-item {counter-increment: ref;} | |
| #refs .gallery-item::after {content: "<Picture " counter(ref) ">"; position: absolute; bottom: 4px; right: 4px; | |
| padding: 0 6px; border-radius: 4px; background: rgba(0, 0, 0, 0.6); color: #fff; font-size: 12px; pointer-events: none;} | |
| #refs .thumbnail-item {pointer-events: none;} | |
| #refs .delete-button {opacity: 1;} | |
| #add-ref {display: none;} | |
| #foot {font-size: 13px; color: var(--body-text-color-subdued);} | |
| """ + f"body:has(#refs .gallery-item):not(:has(#refs .gallery-item:nth-child({MAX_REFS}))) #add-ref {{ display: flex; }}" | |
| BADGE = (HERE / "assets" / "fine-tuned-by-viggle-currentcolor.svg").read_text() | |
| APP = "https://viggle.ai/meme/app/home?ui=v2" | |
| with gr.Blocks(title="Viggle Remake", theme=THEME, css=CSS, head=REFS_HEAD) as demo: | |
| gr.HTML(padding=False, value=f"<div id='hdr'><h1>Viggle Remake</h1><span class='badge'>{BADGE}</span></div>") | |
| gr.Markdown( | |
| "Re-render a video with its people recast (**swap**) or its scene changed (**edit**). The source keeps " | |
| "the motion, camera, framing and soundtrack; the reference pictures show what the result should look like, " | |
| "and with a pin frame each one is held as a keyframe at that frame of the output. 6 sampling steps, " | |
| f"{MIN_UPLOAD_S:.0f}–{MAX_FRAMES/FPS:.1f} s at 24 fps; longer sources are cut.\n\n" | |
| "The prompt is read verbatim and has to follow MiniMax-H3's six-section format " | |
| "(`subject_definitions`, `summary`, `retention_analysis`, `detailed_description`, `overall_soundscape`, " | |
| "`non_diegetic_music`), naming the source `<Video 1>` and the references `<Picture 1>`… in upload order. " | |
| "Its timestamps are output time. Click an example to see one.\n\n" | |
| f"[Model card]({MODEL_CARD}) · **[Viggle app]({APP})**: no setup, no GPU", elem_id="intro") | |
| with gr.Row(equal_height=False): | |
| with gr.Column(): | |
| mode = gr.Radio(["swap", "edit"], value="swap", label="adapter") | |
| v = gr.Video(label="source video") | |
| # type="filepath": prepare() opens the uploads itself. ".webp": the gallery validates by mimetype | |
| refs = gr.Gallery(label=f"reference pictures (1–{MAX_REFS}; drop several at once)", type="filepath", | |
| file_types=["image", ".webp"], interactive=True, columns=3, height="auto", | |
| object_fit="contain", allow_preview=False, show_share_button=False, | |
| show_fullscreen_button=False, elem_id="refs") | |
| add_ref = gr.Button("+ Add picture", size="sm", elem_id="add-ref") | |
| pins = gr.Textbox(label="pin frames, one per picture in order (e.g. 12, 96); empty = no keyframes", | |
| placeholder="12, 96") | |
| prompt = gr.Textbox(lines=16, max_lines=40, label="prompt (six-section format)") | |
| with gr.Row(): | |
| seed = gr.Number(value=0, precision=0, label="seed") | |
| go = gr.Button("Render", variant="primary") | |
| with gr.Column(): | |
| out = gr.Video(label="render") | |
| log = gr.Textbox(label="", lines=3) | |
| gr.Examples( | |
| examples=[[m, str(EX / n / "source.mp4"), (EX / n / "prompt.txt").read_text(), | |
| [str(EX / n / "ref1.png"), str(EX / n / "ref2.png")], ", ".join(map(str, p)), str(EX / n / "out.mp4")] | |
| for n, (m, p, _) in CASES.items()], | |
| inputs=[mode, v, prompt, refs, pins, out], | |
| label="examples — each fills the inputs and shows this app's render of them (seed 0)") | |
| gr.Markdown(f"Powered by [Viggle Remake](https://huggingface.co/Viggle/Viggle-Remake). " | |
| f"Try Viggle's own [Viggle app]({APP}).", elem_id="foot") | |
| add_ref.click(None, js=ADD_REF_JS) | |
| go.click(remake, [mode, v, prompt, refs, pins, seed], [out, log], api_name="remake") | |
| if __name__ == "__main__": | |
| demo.queue(max_size=16).launch(show_error=True, allowed_paths=[str(EX)], max_file_size="200mb") | |