Spaces:
Running on Zero
Running on Zero
| # SPDX-License-Identifier: Apache-2.0 | |
| """Viggle-Animate on ZeroGPU: put the person from a still into a driving video. | |
| Two stages, deliberately split so the OpenAI call never burns GPU quota: | |
| prepare (CPU) frame 0 of the driving clip is handed to gpt-image-2 together with the | |
| character still, and comes back with the performer replaced. That | |
| repainted frame is the model's reference -- the whole "no pose, no | |
| segmentation, no mask" claim rests on this being the only preprocessing. | |
| A frame painted elsewhere can be supplied directly instead, which is what | |
| the bundled examples do: it costs nothing and it renders deterministically. | |
| render (GPU) 4 sampling steps against the driving clip. ~62 GiB of weights, so this | |
| needs ZeroGPU `xlarge` (the full 96 GB card); `large` is 48 GB and | |
| cannot hold it. | |
| Measured on this hardware: weights download at ~1.05 GiB/s and load in ~85 s, for 73.5 GiB | |
| of parameters against a 95.0 GiB card. | |
| The load is at MODULE level because `@spaces.GPU` bodies run in a worker process forked on | |
| the first GPU request -- it inherits everything import time set up, and nothing a request | |
| handler later assigns to a global. The weights stay on the CPU at startup; pre-staging them | |
| through ZeroGPU's tensor pack was tried and does not fit here -- see the note at the bottom | |
| of the load. | |
| """ | |
| import base64 | |
| import contextlib | |
| import inspect | |
| import io | |
| import json | |
| import os | |
| import shutil | |
| import subprocess | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| # 243 frames peaks at 86.1 GiB allocated against a 95.0 GiB card. That margin only survives if | |
| # the allocator can grow segments instead of stranding them, so this is set before torch loads. | |
| # Both spellings: 2.9 renamed the variable and warns on the old one, 2.8 only knows the old one. | |
| os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from huggingface_hub import HfApi, hf_hub_download, snapshot_download | |
| from huggingface_hub.errors import EntryNotFoundError | |
| from openai import OpenAI | |
| from PIL import Image | |
| # A private module: LazyAOTIModel and PACKAGE_FILENAME are not among the package's lazy | |
| # top-level exports. The builder pins `spaces` itself (0.51.1 at this writing -- listing it in | |
| # requirements.txt at any other version is a pip conflict); this module is byte-identical in | |
| # 0.51.1 and 0.51.3, the versions the code was written against. | |
| from spaces.zero.torch.aoti import (PACKAGE_FILENAME, LazyAOTIModel, aoti_capture, | |
| aoti_compile_and_save, aoti_patch) | |
| H3 = "/tmp/h3" | |
| VA = "/tmp/viggle" | |
| MODEL_CARD = "https://huggingface.co/Viggle/Viggle-Animate" | |
| SITE = "https://viggle.ai/h3" | |
| EX = Path(__file__).parent / "examples" | |
| # Every case ships as driving clip + already-painted frame 0 + the render those two produce, | |
| # so clicking one costs no GPU and no OpenAI call. Several share a driving clip, hence the | |
| # name -> (drive file prefix, caption) shape: `corgi` and `clay` reuse `costume`'s clip, and | |
| # the second take of each two-character pair reuses its partner's. | |
| # Order is what a visitor reads top to bottom. `lipsync` is first because it is the only case | |
| # whose driving clip carries a soundtrack, and therefore the only one where the mouth has | |
| # anything to track: every other output's audio track is this finetune's digital silence | |
| # (-105 dBFS, against -22 for this one). Then the swaps where the target is not a person -- one | |
| # paint and one render each, no best-of-N -- and last the human cases the model card compares | |
| # against Wan2.2-Animate. Within that, aspect ratio is grouped, because the output rows below | |
| # are split on these indices. | |
| CASES = { | |
| "lipsync": ("lipsync", "the driving clip has sound, so the mouth has something to say " | |
| "— turn the volume on"), | |
| "corgi": ("costume", "a corgi"), | |
| "panda": ("panda", "a giant panda"), | |
| "penguin": ("penguin", "a penguin"), | |
| "plane": ("plane", "an airliner — wings for arms, landing gear for legs"), | |
| "robot": ("robot", "a walking robot"), | |
| "clay": ("costume", "claymation"), | |
| "duo_anime": ("duo_anime", "two characters at once — cel-shaded cat and corgi"), | |
| "duo_real": ("duo_anime", "the same frame, painted photoreal over the anime set"), | |
| "wushu_animals": ("wushu_animals", "a wide arena shot, both fighters replaced"), | |
| "wushu_robots": ("wushu_animals", "the same shot, gold and chrome robots"), | |
| "prop": ("prop", "a light-up toy, held in both hands"), | |
| "costume": ("costume", "a bow headwrap and floral leggings, driven by a bearded man"), | |
| "style": ("style", "a flat 2D anime character"), | |
| } | |
| FPS = 24 | |
| # The video VAE only encodes frame counts of the form 17n+5, and the pipeline refuses anything | |
| # under 5.0 s, so 124 (5.167 s) is the shortest render that exists -- not a choice made here. | |
| # Uploads are held out to the next such count and the render is cut back, so any length between | |
| # MIN_UPLOAD_S and MAX_FRAMES/FPS comes back at its own length. | |
| # | |
| # The ceiling is VRAM, not the model, which allows up to 345. Measured peak allocated on one | |
| # B200 at 480x832: 79.9 GiB at 124 frames, 82.5 at 175, 86.1 at 243, and 86.8 at 243 frames on a | |
| # 480x864 canvas -- 414,720 px, wider than anything LADDER can now produce, so that 86.8 GiB is a | |
| # true upper bound with 8.2 GiB of headroom against this card's 95.0 GiB. Hence 243, not 345. | |
| CHUNK, TAIL = 17, 5 | |
| MIN_FRAMES, MAX_FRAMES = 124, 243 # 5.167 s .. 10.125 s | |
| MIN_UPLOAD_S = 4.0 | |
| # The 480-class bucket ladder, copied from `ft/buckets.py`. These nine geometries are the only ones | |
| # the model was ever trained or evaluated on. They are 32-aligned by construction and all spend the | |
| # same ~390 rows per latent frame -- the ladder holds *area* constant, not the short edge. | |
| LADDER = [(416, 960), (448, 896), (480, 832), (544, 736), (640, 640), | |
| (736, 544), (832, 480), (896, 448), (960, 416)] | |
| # Measured on this card (renders of 2026-09-01/05): the denoise loop runs 1.6x the B200 timings | |
| # `budget` extrapolates from (42 s at 124 frames), and the host-to-device copy is 22 s from warm | |
| # page cache, 75 s typical, 154 s cold -- and one 261 s outlier that overran the request and drew | |
| # only a warning. Both are set high on purpose -- the padded request is what locks the free tier | |
| # out at admission; see `budget`. | |
| SLOWDOWN = float(os.environ.get("SLOWDOWN", "1.6")) # this card vs the B200 the timings came from | |
| H2D = float(os.environ.get("H2D", "170")) # seconds to move the weights onto the card | |
| # Measured good on this card: the first render logged no fallback, so sm_120 serves the shape. | |
| ATTN = os.environ.get("ATTN", "_native_cudnn") # set to "" to stay on torch's default | |
| DAILY = int(os.environ.get("DAILY_LIMIT", "5")) | |
| # Who may press the AOTI build button, which holds a write token for the weights repo. | |
| OWNERS = {u.strip().lower() for u in os.environ.get("OWNERS", "yycc").split(",") if u.strip()} | |
| # The repaint is downscaled to a ladder canvas -- 480x832 at most -- before the model ever sees | |
| # it, so the detail the higher tiers buy is resampled away. What has to survive that downscale is | |
| # the pose and the identity, not the texture. "low", "medium", "high" or "auto". | |
| QUALITY = os.environ.get("IMAGE_QUALITY", "low") | |
| # There is no `input_fidelity` here on purpose. The openai SDK accepts the parameter, but | |
| # gpt-image-2 rejects it -- `invalid_input_fidelity_model`, a 400 on every request, not a | |
| # warning. It is a gpt-image-1 knob. Do not add it back. | |
| FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" | |
| # Fail at boot, not in `finish` -- which runs after the GPU money is already spent. | |
| assert Path(FONT).exists(), f"watermark font missing: {FONT}" | |
| t = time.time() | |
| snapshot_download("MiniMaxAI/MiniMax-H3", local_dir=H3, | |
| allow_patterns=["*.json", "vae/*", "audio_vae/*", | |
| "scheduler/*", "audio_scheduler/*"]) | |
| snapshot_download("Viggle/Viggle-Animate", local_dir=VA, | |
| allow_patterns=["transformer/*", "lora/*", "assets/*"]) | |
| print(f"[boot] weights in {time.time()-t:.0f}s", 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 MiniMaxH3SetTimestepsStep | |
| from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3VideoDecodeStep | |
| from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3LoopSchedulerStep | |
| from diffusers.modular_pipelines.minimax_h3.encoders import MiniMaxH3Ref2VATextEncoderStep | |
| from diffusers.modular_pipelines.minimax_h3.modular_pipeline import audio_latent_num_frames | |
| from diffusers.utils.export_utils import encode_video | |
| # The text encoder is never loaded: conditioning is a frozen 362x5120 tensor computed once | |
| # from a fixed prompt, so Qwen3-VL (63 GB of the base repo) stays on disk. Nothing in the | |
| # output comes from text anyone types here -- there is no prompt box for that reason. | |
| FIXED = torch.load(f"{VA}/assets/fixed_embed_fwd_anyframe.pt", weights_only=False) | |
| def use_fixed_embeds(self, components, state): | |
| block_state = self.get_block_state(state) | |
| block_state.prompt_embeds = FIXED["prompt_embeds"].to(components._execution_device, | |
| torch.bfloat16) | |
| block_state.text_token_tags = FIXED["text_token_tags"] | |
| self.set_block_state(state, block_state) | |
| return components, state | |
| MiniMaxH3Ref2VATextEncoderStep.__call__ = use_fixed_embeds | |
| # stock has it too; the pipeline already runs under no_grad, so belt-and-braces | |
| def decode_to_uint8(self, components, state): | |
| """The stock step ends in `postprocess_video`, which walks the decoded clip to fp32 numpy on | |
| the CPU -- measured 9.2 s at 243 frames against 0.03 s for quantizing on the card | |
| (REALTIME_PIPELINE.md: "pure format conversion, buys nothing"). Everything through the fp32 | |
| denormalize below is the stock code, and fp32 round is round-half-to-even in numpy and torch | |
| alike, so the u8 frames are bit-identical to what the old path fed x264 -- only the CPU round | |
| trip is gone. `encode_video` takes the [T,H,W,C] uint8 tensor as-is. `output_type` is | |
| ignored: this app is the only caller and it always wants frames. | |
| """ | |
| 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) | |
| # In-place on purpose: `video` is a fresh tensor never read again, and the out-of-place | |
| # spelling would hold two extra full-size fp32 temporaries at the peak. | |
| 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 | |
| # --------------------------------------------------------------------------------------------- | |
| # The audio pin. When the upload has a soundtrack it is VAE-encoded and written into the render's | |
| # *target* audio rows as a clean latent -- t=1, H3's flow being reversed -- and held there for the | |
| # whole loop, instead of being denoised out of noise like every other target row. The video branch | |
| # attends to those rows, so the picture is driven by the real speech. That is what improves | |
| # lip-sync, and it is measurable rather than hoped for: against a control-vs-control floor of | |
| # 50.4 dB, swapping the pinned track for encoded silence moves the video by 23.0-32.8 dB across | |
| # four shot/canvas combinations. It costs nothing -- peak VRAM matched the control to 0.1 GiB and | |
| # wall clock to within run-to-run noise. `AUDIO_PIN.md` in the repo has the falsification runs. | |
| # | |
| # Both patches are installed once, at import, and read PIN at call time. Installing them per | |
| # request would leak across visitors: a class patch is process-global and one worker serves every | |
| # request, so an upload that had audio would silently pin the next upload that did not. | |
| PIN = False | |
| _build_row_timesteps = MiniMaxH3SetTimestepsStep.build_row_timesteps | |
| def pinned_row_timesteps(*a): | |
| """Tell the transformer the target audio rows are already clean. | |
| Positional arg 6 is `audio_timestep`, and `before_denoise.py` passes all nine positionally. | |
| Without this the rows hold real audio while the model is told they are noise at the current | |
| sigma, and it denoises against a premise that is false. | |
| """ | |
| return _build_row_timesteps(*a[:6], 1.0, *a[7:]) if PIN else _build_row_timesteps(*a) | |
| MiniMaxH3SetTimestepsStep.build_row_timesteps = staticmethod(pinned_row_timesteps) | |
| _loop_step = MiniMaxH3LoopSchedulerStep.__call__ | |
| def video_only_step(self, components, block_state, i, t): | |
| """Step the video rows only, so the pinned audio rows survive every step untouched. | |
| The obvious alternative -- raising `num_condition_audio_rows` so the loop treats them as | |
| conditioning -- does not work: it empties the decoder's slice | |
| `audio_latents[num_condition_audio_rows:]` and trips the reference-count assert in | |
| `MiniMaxH3PrepareLatentsStep`. That count stays 0. | |
| """ | |
| if not PIN: | |
| return _loop_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 | |
| t = time.time() | |
| PIPE = ModularPipeline.from_pretrained(H3, workflow="ref2va") | |
| PIPE.load_components(names=["vae", "audio_vae", "scheduler", "audio_scheduler"], | |
| pretrained_model_name_or_path=H3, torch_dtype=torch.bfloat16) | |
| PIPE.transformer_ref = MiniMaxH3Transformer3DModel.from_pretrained( | |
| f"{VA}/transformer", torch_dtype=torch.bfloat16) | |
| # `prefix=None` and the explicit `weight_name` are both required: the loader defaults to a | |
| # `.bin` name (raises) and to filtering keys for a `transformer.` prefix these bare keys do | |
| # not have -- that mismatch loads nothing and only warns, i.e. silently renders the | |
| # un-distilled model at 4 steps, which looks like noise. | |
| PIPE.transformer_ref.load_lora_adapter(f"{VA}/lora", | |
| weight_name="pytorch_lora_weights.safetensors", | |
| prefix=None) | |
| PIPE.scheduler.set_shift(3.0) # the base default is 12; the few-step student wants 3 | |
| ASR = PIPE.audio_vae.config.sampling_rate # 32 kHz; the rate the pin has to be fed at | |
| # AOT-compiled copies of the transformer's two repeated block classes, built ON THIS HARDWARE by | |
| # the maintenance button below (a package compiled on another arch will not load) and stored in | |
| # the weights repo as <ClassName>/package.pt2 -- code only, tens of MB: the weights come from | |
| # each block's own state_dict at patch time. Missing packages mean an eager Space, not a broken | |
| # one. Only the *download* happens here; the patching waits for the GPU worker, because | |
| # `aoti_patch` snapshots `state_dict()` and until `PIPE.to("cuda")` in `render` those storages | |
| # are the CPU ones the copy abandons -- patched at boot, the compiled blocks would read CPU | |
| # weights forever. See `wire_aoti`. | |
| AOTI_REPO = "Viggle/Viggle-Animate" | |
| class _Cell: | |
| # Stand-in for the ContextVars inside LazyAOTIModel. The worker runs every GPU task in a | |
| # fresh copy_context(), so a ContextVar set during one render is invisible to the next -- | |
| # the dlopen'd package would be re-extracted and re-loaded inside every billed window, and | |
| # each extraction is another copy charged against the container's hidden write quota. A | |
| # plain attribute survives across tasks; the parent never calls these models, so the | |
| # context isolation being discarded here protected nothing. | |
| def __init__(self): | |
| self.v = None | |
| def get(self): | |
| return self.v | |
| def set(self, v): | |
| self.v = v | |
| AOTI_BLOCKS = {} | |
| # `ATTN=""` must keep meaning "no cuDNN anywhere": the packages baked cuDNN in at build time | |
| # and ignore `set_attention_backend`, so with the kill switch set they may not be wired at all. | |
| if os.environ.get("AOTI", "1") == "1" and ATTN: | |
| try: | |
| AOTI_BLOCKS = {name: LazyAOTIModel(hf_hub_download(AOTI_REPO, PACKAGE_FILENAME, | |
| subfolder=name)) | |
| for name in MiniMaxH3Transformer3DModel._repeated_blocks} | |
| for lazy in AOTI_BLOCKS.values(): | |
| lazy.compiled_model, lazy.loaded_weights = _Cell(), _Cell() | |
| print(f"[boot] aoti packages loaded: {', '.join(AOTI_BLOCKS)}", flush=True) | |
| except EntryNotFoundError: | |
| AOTI_BLOCKS = {} | |
| print("[boot] no aoti packages on the hub -- rendering eager", flush=True) | |
| # The weights stay on the CPU until `render`'s `PIPE.to("cuda")`. Pre-staging them through | |
| # ZeroGPU's tensor pack (`PIPE.to("cuda")` here at import) was tried 2026-09-01 and cannot fit: | |
| # the pack `posix_fallocate`s 78.6 GB while the 80 GB of sources still sit on the same disk, and | |
| # the container has a hidden ~100 GB write quota that statvfs does not report -- the boot line | |
| # below prints >1.3 TB free while the fallocate dies ENOSPC instantly. The prize was small | |
| # anyway: the pack loads via O_DIRECT disk reads, the copy it replaces runs from warm page cache. | |
| print(f"[boot] pipeline in {time.time()-t:.0f}s, " | |
| f"disk free {shutil.disk_usage('/tmp').free/1e9:.0f} GB /tmp, " | |
| f"{shutil.disk_usage(Path.home()).free/1e9:.0f} GB {Path.home()}", flush=True) | |
| def run(*cmd): | |
| # ffmpeg says why it failed on stderr, and `capture_output` swallows it -- a bare | |
| # CalledProcessError names the exit code and nothing else, which is useless for a codec or | |
| # container the uploader chose. | |
| p = subprocess.run(cmd, capture_output=True, text=True, errors="replace") | |
| if p.returncode: | |
| raise gr.Error(f"{cmd[0]} failed: {p.stderr.strip().splitlines()[-1] if p.stderr.strip() else p.returncode}") | |
| def sweep(hours=6): | |
| """Drop working directories older than `hours`. | |
| Each request leaves one behind holding a few hundred MB of intermediate video, and the | |
| container runs for weeks. Pruning by age rather than tracking them is deliberate: a render | |
| in flight still owns its directory, and age is the one property that cannot mistake a live | |
| job for a dead one. | |
| """ | |
| cutoff = time.time() - hours * 3600 | |
| for d in Path(tempfile.gettempdir()).glob("tmp*"): | |
| # Two concurrent requests both sweep, so a directory can disappear between the glob and | |
| # the stat. That is not a reason to fail somebody's render. | |
| with contextlib.suppress(OSError): | |
| if d.is_dir() and d.stat().st_mtime < cutoff: | |
| shutil.rmtree(d, ignore_errors=True) | |
| def probe(path): | |
| out = subprocess.run( | |
| ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", | |
| "stream=width,height", "-of", "json", path], | |
| check=True, capture_output=True, text=True).stdout | |
| s = json.loads(out)["streams"][0] | |
| return s["width"], s["height"] | |
| def nframes(path): | |
| return int(subprocess.run( | |
| ["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames", | |
| "-show_entries", "stream=nb_read_frames", "-of", "csv=p=0", path], | |
| check=True, capture_output=True, text=True).stdout.strip()) | |
| def has_audio(path): | |
| """Whether the upload carries a soundtrack, which is what decides if the pin is used.""" | |
| return bool(subprocess.run( | |
| ["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries", | |
| "stream=codec_type", "-of", "csv=p=0", path], | |
| check=True, capture_output=True, text=True).stdout.strip()) | |
| def grid(n): | |
| """Smallest frame count the video VAE accepts (17k+5) that still covers `n` frames. | |
| Rounding *up* costs at most 16 frames of held-still tail, which gets trimmed off the render. | |
| Rounding down would instead silently shorten every clip that is not already on the grid. | |
| """ | |
| return min(MAX_FRAMES, max(MIN_FRAMES, -(-(n - TAIL) // CHUNK) * CHUNK + TAIL)) | |
| def canvas(w, h): | |
| """The ladder geometry whose aspect is closest to the upload's. | |
| Not "scale to a 480 short edge", which is the obvious rule and the wrong one. The ladder | |
| holds area constant, not the short edge, so the two rules agree only near 480x832 and drift | |
| apart everywhere else -- a square clip on a 480 short edge is 480x480, which is 225 rows per | |
| latent frame against the ladder's 390. `ft/dataset_fwd.py:59` records training doing exactly | |
| that and undoing it: "1:1 used to resize to 480x480 [...] 42% of the budget thrown away on | |
| the one class that could most afford the pixels. Index 4 is 640x640." | |
| Snapping to the ladder also removes two failure modes for free. Every entry is a multiple of | |
| 32, so the patchify guard can no longer fire -- it is what the 16-grid 464 case used to trip. | |
| And the aspects are bounded to 0.43..2.31, so a panorama or a column no longer produces a | |
| canvas the model has never seen; it gets cover-cropped to the nearest one it has. | |
| """ | |
| return min(LADDER, key=lambda b: abs(b[0] / b[1] - w / h)) | |
| def fit(img, w, h): | |
| """Land a reference frame on exactly w x h without stretching it. | |
| Stretching would be the quiet failure here: the reference is compared against the driving | |
| clip frame by frame, so a horizontally squeezed body is a body that disagrees with every | |
| frame of the motion. Off-aspect input is cover-cropped instead, losing edges rather than | |
| proportions. | |
| """ | |
| if abs(img.width / img.height - w / h) > 0.01: | |
| s = max(w / img.width, h / img.height) | |
| img = img.resize((round(img.width * s), round(img.height * s)), Image.LANCZOS) | |
| img = img.crop(((img.width - w) // 2, (img.height - h) // 2, | |
| (img.width - w) // 2 + w, (img.height - h) // 2 + h)) | |
| return img.resize((w, h), Image.LANCZOS) | |
| # The repainted frame is a motion anchor before it is a picture. Whatever geometry it shows is | |
| # what the video model tries to reconcile with the driving clip, and when the two disagree the | |
| # driving clip wins -- an editor that quietly re-poses or re-centres the subject produces a | |
| # reference that fights every frame. So the pose is stated as the constraint and the identity as | |
| # the single edit, rather than both being items in one list. | |
| REPAINT = ( | |
| "Edit the first image, which is one frame of a video. Change exactly one thing: who the " | |
| "person is. Give them the face, hair, skin tone and clothing of the person in the second " | |
| "image.\n" | |
| "The body does not move. Keep the pose identical -- every limb, both hands, the head angle " | |
| "and the direction of gaze stay exactly where they are, and the person keeps the same " | |
| "position in the frame and the same size. Do not re-pose, straighten, re-centre or redraw " | |
| "the body: repaint the person who is already standing there, do not draw a new one.\n" | |
| "Nothing else changes. Same camera angle, same framing and crop, same background, same " | |
| "lighting, shadows and colour grade. Anything the person is holding or touching stays " | |
| "exactly as it is, in the same hand, at the same angle.\n" | |
| "Do not zoom or crop, and do not add text, captions or watermarks." | |
| ) | |
| SEEN = {} | |
| def gate(request): | |
| """One repaint counter per client IP per day. | |
| This is a spend limit on the owner's image-model budget, not on the GPU. The gpt-image-2 | |
| repaint happens in `prepare`, on the CPU, before any `@spaces.GPU` call -- ZeroGPU's quota | |
| never sees it, and only this counter stops it. The GPU needs no gate here at all: ZeroGPU's | |
| own admission check turns visitors away by their Hub quota (see `budget`). | |
| An IP is a weaker key than the Hub username this used to count -- a forwarded header can be | |
| rewritten, and one NAT can hide a classroom -- but the username needed a sign-in, and the | |
| sign-in was judged to cost more visitors than the key was worth. A brake, not a wall. | |
| """ | |
| key = _client_key(request) | |
| used = SEEN.get(key, 0) | |
| if used >= DAILY: | |
| raise gr.Error(f"daily repaint limit reached ({DAILY}). Try again tomorrow, or supply " | |
| "a first frame you painted elsewhere -- that path is not limited.") | |
| # Yesterday's counters are dead weight in a process that stays up for weeks. | |
| for k in [k for k in SEEN if k[1] != key[1]]: | |
| del SEEN[k] | |
| SEEN[key] = used + 1 | |
| return DAILY - used - 1 | |
| def _client_key(request): | |
| ip = (request.headers.get("x-forwarded-for", "").split(",")[0].strip() | |
| or request.client.host) if request else "?" | |
| return ip, time.strftime("%Y-%m-%d") | |
| def refund(request): | |
| # A repaint the server failed to deliver should not cost one of the day's five. | |
| key = _client_key(request) | |
| if SEEN.get(key, 0) > 0: | |
| SEEN[key] -= 1 | |
| # A Space restart empties /tmp/gradio, but pages opened before it still reference files there. | |
| # The re-submitted click then dies inside the components' own preprocess (PIL.Image.open -- | |
| # production, 2026-09-05), *before* any handler runs, so only a wrapper here can turn the raw | |
| # FileNotFoundError into something a visitor can act on. | |
| _GONE = ("this page predates a Space restart and its uploads are gone -- " | |
| "reload the page and upload the files again") | |
| for _cls in (gr.Image, gr.Video): | |
| def _guarded(self, *a, _orig=_cls.preprocess, **k): | |
| try: | |
| return _orig(self, *a, **k) | |
| except FileNotFoundError: | |
| raise gr.Error(_GONE) | |
| _cls.preprocess = _guarded | |
| def prepare(video, character, painted, request: gr.Request = None): | |
| # `request` is not wired to any component: gradio fills it in from the type annotation. | |
| if not video: | |
| raise gr.Error("upload a driving video") | |
| if not character and not painted: | |
| raise gr.Error("give either a character image to paint in, or a first frame you have " | |
| "already painted yourself") | |
| # Backstop for the preprocess guard below: covers any stale path that slips through as a | |
| # bare filepath without gradio touching the file first. | |
| if any(p and not os.path.exists(p) for p in (video, character, painted)): | |
| raise gr.Error(_GONE) | |
| sweep() | |
| work = Path(tempfile.mkdtemp()) | |
| sw, sh = probe(video) | |
| w, h = canvas(sw, sh) | |
| # For any ordinary aspect the nearest ladder entry is within ~3%, which is a trim nobody | |
| # notices. A panorama or a tall column is a different matter: it is far from every entry, so | |
| # cover-cropping to one throws away most of the frame. Say so rather than let someone wonder | |
| # where their shot went -- the alternative, refusing, would be rendering nothing at all. | |
| lost = 1 - min((sw / sh) / (w / h), (w / h) / (sw / sh)) | |
| note = f" — cropped {lost*100:.0f}% to reach it" if lost > 0.15 else "" | |
| # Cover-crop, not stretch, for the same reason `fit` does it to the reference: the two are | |
| # compared frame by frame, so they have to be squeezed the same way or not at all. `canvas` | |
| # picks the nearest ladder aspect, so what is cropped here is the few percent by which the | |
| # upload misses it -- unless the upload is a panorama, which loses its edges by design. | |
| 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), "-crf", "6", "-an", full) | |
| have = nframes(full) | |
| if have < round(MIN_UPLOAD_S * FPS): | |
| raise gr.Error(f"driving video is {have/FPS:.1f}s after retiming to {FPS} fps; " | |
| f"at least {MIN_UPLOAD_S:.1f}s is needed") | |
| # The model renders on the 17k+5 grid and never shorter than 124 frames, so the clip is | |
| # held out to the next legal count with a frozen last frame and the render is cut back to | |
| # `keep` afterwards. That padding is what makes a 4 s upload possible at all. | |
| frames = grid(have) | |
| 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), "-crf", "6", "-an", cond) | |
| keep = have / FPS | |
| # The upload's own soundtrack, kept for the pin. Both ffmpeg calls above still pass `-an`: | |
| # the driving clip is a *reference*, and its audio must not become conditioning. This is a | |
| # separate extraction for a different destination -- the render's target audio rows. | |
| # | |
| # Cut to `keep`, not to the padded `frames`: the tail the 17k+5 grid forces us to render is a | |
| # frozen frame, and real audio under a frozen frame is exactly the desync the pin exists to | |
| # remove. The rest is zero-padded at encode time. Raw f32 rather than a container so both | |
| # consumers read it without a decoder -- `render` with torch.frombuffer, `finish` with -f f32le. | |
| 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) | |
| ref = str(work / "ref.png") | |
| # Already painted: there is nothing to generate, only a canvas to agree on. This is the | |
| # path the bundled examples take -- no OpenAI call, and the render is then deterministic | |
| # in the seed, which the gpt-image-2 path can never be. | |
| if painted: | |
| fit(Image.open(painted).convert("RGB"), w, h).save(ref) | |
| return (ref, cond, keep, frames, apath, ref, None, | |
| f"canvas {w}x{h}{note}, {keep:.1f}s out{', audio pinned' if apath else ''}") | |
| # The counter is taken here, not at the top: it exists for the OpenAI spend below, so the | |
| # painted path is uncounted and a rejected upload never costs one of the day's five. | |
| left = gate(request) | |
| # `OpenAI()` picks up OPENAI_API_KEY *and* OPENAI_BASE_URL from the environment, so which | |
| # provider serves the repaint is a Space-secret decision, not a code one. The two have to | |
| # agree: a relay key against the default endpoint is a 401 on every request. | |
| if not os.environ.get("OPENAI_API_KEY"): | |
| raise gr.Error("OPENAI_API_KEY is not set on this Space") | |
| frame0 = str(work / "frame0.png") | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", cond, "-frames:v", "1", frame0) | |
| # gpt-image-2 accepts exactly three sizes -- anything else is a 400 whose generic message | |
| # happens to contain "size", which made every failure here look like an aspect problem | |
| # (production, 2026-09-05). So: canonical orientation straight away, cover-crop back onto | |
| # the canvas afterwards. | |
| size = "1024x1536" if h > w else "1536x1024" if w > h else "1024x1024" | |
| t0 = time.time() | |
| b64 = None | |
| for attempt in range(3): | |
| try: | |
| with open(frame0, "rb") as a, open(character, "rb") as b: | |
| r = OpenAI(timeout=300).images.edit( | |
| model="gpt-image-2", image=[a, b], prompt=REPAINT, size=size, | |
| quality=QUALITY, output_format="png", n=1) | |
| b64 = r.data[0].b64_json | |
| break | |
| except Exception as e: | |
| print(f"[paint] attempt {attempt}: {e}", flush=True) | |
| if "moderation_blocked" in str(e): | |
| # Deterministic: the same images meet the same safety system, so retrying is | |
| # three refusals for the price of three. It still costs one of the day's five; | |
| # that counter is the only brake on hammering the filter with variations. | |
| raise gr.Error("OpenAI's safety system refused this repaint. Photoreal people " | |
| "trip it most often; stylized or costumed characters usually " | |
| "pass.") | |
| if attempt == 2: | |
| low = str(e).lower() | |
| if "server_error" in low or "connection" in low or "timed out" in low: | |
| # The visitor did nothing wrong; during the 09-03 relay outage every 504 | |
| # still burned quota. Give it back. | |
| refund(request) | |
| raise gr.Error(f"gpt-image-2 failed: {e}") | |
| time.sleep(3) | |
| img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB") | |
| fit(img, w, h).save(ref) # the canonical size is never the canvas aspect; cover-crop | |
| return (ref, cond, keep, frames, apath, ref, None, | |
| f"repainted in {time.time()-t0:.0f}s, canvas {w}x{h}{note}, {keep:.1f}s out" | |
| f"{', audio pinned' if apath else ''}, {left} repaints left today") | |
| _WIRED = False | |
| def wire_aoti(): | |
| """Swap each repeated block's forward for its compiled package -- in the GPU worker, *after* | |
| `PIPE.to("cuda")`. `aoti_patch` snapshots `state_dict()` and holds whatever storages the | |
| params have at that moment, so patching before the copy would feed the compiled kernels the | |
| CPU weights the copy abandons. Latched per worker: the parent process never wires (it has no | |
| GPU and `render` never runs there), each forked worker wires on its first render and stays | |
| wired while warm. One caveat rides along: the compiled graph baked the attention backend | |
| that was active at *build* time (cudnn), so `set_attention_backend` no longer reaches these | |
| blocks -- the cuDNN retry in `render` has to unwire before it means anything. | |
| """ | |
| global _WIRED | |
| if _WIRED or not AOTI_BLOCKS: | |
| return | |
| n = 0 | |
| for name, lazy in AOTI_BLOCKS.items(): | |
| for m in PIPE.transformer_ref.modules(): | |
| if type(m).__name__ == name: | |
| aoti_patch(m, lazy) | |
| n += 1 | |
| _WIRED = True | |
| print(f"[gpu] aoti wired {n} blocks", flush=True) | |
| def unwire_aoti(): | |
| """Back to eager, permanently for this worker: `aoti_patch` writes an instance attribute | |
| that shadows the class `forward`, so popping it is a complete undo. Reached from `render`'s | |
| retry, which catches *any* first-attempt failure -- an OOM at 243 frames as much as a cuDNN | |
| refusal -- and staying eager for good is the conservative answer to every one of those | |
| causes.""" | |
| global AOTI_BLOCKS, _WIRED | |
| for m in PIPE.transformer_ref.modules(): | |
| if type(m).__name__ in AOTI_BLOCKS: | |
| m.__dict__.pop("forward", None) | |
| AOTI_BLOCKS = {} | |
| _WIRED = False | |
| # The geometry the packages are built at, and the different one they are validated at. Every | |
| # bundled example is 480x832, so the second geometry -- the one that proves the sequence length | |
| # stayed a dynamic dim instead of silently specializing -- has to be synthesized (`build_aoti` | |
| # cover-crops an example drive to it). 544x736 is a real ladder entry, so it is also a shape | |
| # production will actually serve. | |
| BUILD_CANVASES = [(480, 832), (544, 736)] | |
| # Off by default: max_autotune benchmarks kernel variants per compile and can add tens of | |
| # minutes, and the build has to fit one PRO day's xlarge admission (see build_aoti). | |
| AUTOTUNE = os.environ.get("AOTI_AUTOTUNE", "0") == "1" | |
| def _capture_kwargs(cond, ref, w, h): | |
| # Exactly what `render` passes -- the captured tensors must have the shapes production sees. | |
| # `render` sets PIN on every call, so this cannot inherit a stale True from a render that | |
| # raised; cleared here anyway because this path calls the pipeline without going through it. | |
| global PIN | |
| PIN = False | |
| PIPE.register_to_config(canvas_short_edge=min(w, h), | |
| canvas_max_pixels=min(w, h) * max(w, h), | |
| reference_image_short_edge=min(w, h)) | |
| video = MiniMaxH3VideoReference.from_file(cond) | |
| video.audio, video.sample_rate = None, None | |
| return dict( | |
| prompt=FIXED["presentation"], | |
| references=[video, MiniMaxH3ImageReference.from_file(ref)], | |
| num_frames=MIN_FRAMES, height=h, width=w, num_inference_steps=4, | |
| generator=torch.Generator().manual_seed(42), | |
| output=["videos", "audio", "sampling_rate"], | |
| ) | |
| def build_aoti(): | |
| """Compile the two repeated block classes on this card and hand the packages back as bytes. | |
| One export covers all 52 block instances and every shape bucket: the sequence dims are left | |
| to Dim.AUTO, and the artifact is then run on a geometry the export never saw -- a dim that | |
| silently specialized fails that call loudly instead of failing some visitor's render later. | |
| Inputs come from `aoti_capture`, which aborts a real pipeline call at the block's first | |
| forward -- diffusers logs each abort as a big ERROR traceback; in this build that is the | |
| plan working, not a failure. | |
| Billed like any render, from the clicker's quota -- and admission doubles the request on | |
| xlarge, so duration=1100 asks for 2200 s of a PRO day's 2400: the largest round request | |
| that still clears a fresh day. Only actual elapsed is charged; without autotune the whole | |
| build (4 captures, 2 exports, 2 compiles, a 243-frame probe) should run well under the | |
| ceiling. If it ever hits it, raise AOTI_BUILD_BUDGET -- but past 1200 a PRO account can no | |
| longer be admitted at all. | |
| """ | |
| from torch.export import Dim, export | |
| from torch.utils._pytree import tree_leaves, tree_map | |
| logs = [] | |
| def say(s): | |
| logs.append(s) | |
| print(f"[build] {s}", flush=True) | |
| PIPE.to("cuda") | |
| if ATTN: | |
| PIPE.transformer_ref.set_attention_backend(ATTN) | |
| names = MiniMaxH3Transformer3DModel._repeated_blocks | |
| # An earlier build click on this same worker may have left blocks instance-patched (the | |
| # wired probe below, or a validation that died between patch and unpatch). The export must | |
| # trace eager, and `aoti_capture`'s mock would restore that stale forward instead of the | |
| # class one -- so pop before anything else. Scoped to the block classes: a blanket pop | |
| # would strip forwards diffusers hooks may have installed on other modules. | |
| for m in PIPE.transformer_ref.modules(): | |
| if type(m).__name__ in names: | |
| m.__dict__.pop("forward", None) | |
| firsts = {} | |
| for name in names: | |
| blocks = [m for m in PIPE.transformer_ref.modules() if type(m).__name__ == name] | |
| # Block 0 stands in for the whole class, so every sibling must have the same structure. | |
| # "Same class" stopped guaranteeing that when the LoRA was injected: peft rewrites the | |
| # Linears of the layers it targets, and an adapter that skipped some blocks would make | |
| # their state_dicts (and so the compiled graph's constants) disagree. | |
| assert all(set(b.state_dict()) == set(blocks[0].state_dict()) for b in blocks[1:]), \ | |
| f"{name}: blocks differ in structure; one export cannot cover them" | |
| firsts[name] = blocks[0] | |
| say(f"{name}: {len(blocks)} instances, structurally uniform") | |
| work = Path(tempfile.mkdtemp()) | |
| drive, still = str(EX / "prop_drive.mp4"), str(EX / "prop_ref.png") | |
| cases = [] | |
| for w, h in BUILD_CANVASES: | |
| cond, ref = str(work / f"cond{w}.mp4"), str(work / f"ref{w}.png") | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", drive, "-vf", | |
| f"scale={w}:{h}:force_original_aspect_ratio=increase,crop={w}:{h}", | |
| "-frames:v", str(MIN_FRAMES), "-crf", "6", "-an", cond) | |
| fit(Image.open(still).convert("RGB"), w, h).save(ref) | |
| cases.append((cond, ref, w, h)) | |
| calls = {name: [] for name in names} | |
| for cond, ref, w, h in cases: | |
| kw = _capture_kwargs(cond, ref, w, h) | |
| for name in names: | |
| t0 = time.time() | |
| with aoti_capture(firsts[name]) as call: | |
| PIPE(**kw) | |
| assert call.args or call.kwargs, f"{name} was never called at {w}x{h}" | |
| calls[name].append(call) | |
| say(f"captured {name} at {w}x{h} in {time.time() - t0:.0f}s") | |
| packages = {} | |
| for name in names: | |
| block, (call_a, call_b) = firsts[name], calls[name] | |
| # dynamic_shapes is keyed by the forward's parameter names; bind() recovers them no | |
| # matter how the transformer spells the call. No apply_defaults: a key for a parameter | |
| # the call never passed is an export error, not a wildcard. | |
| bound = inspect.signature(block.forward).bind(*call_a.args, **call_a.kwargs) | |
| dyn = {k: tree_map(lambda v: [Dim.AUTO] * v.ndim if isinstance(v, torch.Tensor) else None, x) | |
| for k, x in bound.arguments.items()} | |
| t0 = time.time() | |
| ep = export(block, call_a.args, call_a.kwargs, dynamic_shapes=dyn) | |
| say(f"exported {name} in {time.time() - t0:.0f}s") | |
| t0 = time.time() | |
| aoti_compile_and_save(work / name, ep, {"max_autotune": AUTOTUNE}) | |
| pkg = work / name / "root" / PACKAGE_FILENAME | |
| say(f"compiled {name} in {time.time() - t0:.0f}s -> {pkg.stat().st_size / 2**20:.0f} MB") | |
| # Run the artifact on the geometry the export never saw, against eager. This is the | |
| # specialization check, and -- because it goes through the same aoti_patch/load path the | |
| # render uses, with check_full_update -- the constants-fqn check for the peft-wrapped | |
| # structure too. The refiner's inputs are the same shape in both cases (its tokens are | |
| # the frozen text, not the video), so for it this is only the fqn half -- and also why | |
| # its specializing to that shape is harmless. | |
| with torch.no_grad(): | |
| want = block(*call_b.args, **call_b.kwargs) | |
| aoti_patch(block, LazyAOTIModel(str(pkg))) | |
| try: | |
| got = block(*call_b.args, **call_b.kwargs) | |
| finally: | |
| del block.forward # even on failure -- the worker survives this fn | |
| want, got = tree_leaves(want), tree_leaves(got) | |
| assert len(want) == len(got), f"{name}: output arity {len(got)} vs {len(want)}" | |
| rel = max((a - b).abs().max().item() / max(a.abs().max().item(), 1e-6) | |
| for a, b in zip(want, got)) | |
| assert rel < 0.05, f"{name}: compiled vs eager rel diff {rel:.3f} -- wrong weights?" | |
| say(f"validated {name} at {cases[1][2]}x{cases[1][3]}: rel diff {rel:.4f}") | |
| packages[name] = pkg.read_bytes() | |
| # The worst case production can serve -- 243 frames at the largest-area canvas -- has only | |
| # ever been measured eager, where it peaks ~86.8 of the card's 95 GiB. Run it once with the | |
| # fresh packages wired exactly as `render` will wire them, so an out-of-memory (dlopen'd | |
| # kernels + their workspaces on top of that peak) fails here, on the owner's click, and not | |
| # on a visitor's longest upload. Also the third geometry: any frame-count-dependent dim | |
| # that both 124-frame build cases shared, and the dual-geometry check above therefore | |
| # could not see, breaks this call instead. | |
| cond, ref = str(work / "probe.mp4"), str(work / "probe.png") | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", drive, "-vf", | |
| "scale=640:640:force_original_aspect_ratio=increase,crop=640:640," | |
| "tpad=stop_mode=clone:stop=-1", | |
| "-frames:v", str(MAX_FRAMES), "-crf", "6", "-an", cond) | |
| fit(Image.open(still).convert("RGB"), 640, 640).save(ref) | |
| lazies = {name: LazyAOTIModel(str(work / name / "root" / PACKAGE_FILENAME)) | |
| for name in names} | |
| try: | |
| for m in PIPE.transformer_ref.modules(): | |
| if type(m).__name__ in lazies: | |
| aoti_patch(m, lazies[type(m).__name__]) | |
| # 2 steps, not 1: the scheduler refuses num_inference_steps < 2 without an explicit | |
| # sigmas schedule (this killed the first two builds, after everything had compiled). | |
| # Peak memory is the same either way. | |
| kw = _capture_kwargs(cond, ref, 640, 640) | dict(num_frames=MAX_FRAMES, | |
| num_inference_steps=2) | |
| torch.cuda.reset_peak_memory_stats() | |
| t0 = time.time() | |
| PIPE(**kw) | |
| say(f"wired probe: {MAX_FRAMES} frames at 640x640 in {time.time() - t0:.0f}s, " | |
| f"peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB") | |
| finally: | |
| for m in PIPE.transformer_ref.modules(): | |
| if type(m).__name__ in lazies: | |
| m.__dict__.pop("forward", None) | |
| return packages, "\n".join(logs) | |
| def build_ui(profile: gr.OAuthProfile | None = None): | |
| # The build burns most of the clicker's daily xlarge quota and writes to the weights repo; | |
| # nobody but the owner should be able to press it even if AOTI_BUILD is left on by mistake. | |
| if profile is None or profile.username.lower() not in OWNERS: | |
| raise gr.Error("owner-only maintenance button") | |
| packages, log = build_aoti() | |
| for name, blob in packages.items(): | |
| # HF_TOKEN comes from the Space secrets via the environment; HfApi picks it up itself. | |
| HfApi().upload_file(path_or_fileobj=io.BytesIO(blob), | |
| path_in_repo=f"{name}/{PACKAGE_FILENAME}", | |
| repo_id=AOTI_REPO, repo_type="model") | |
| return log + "\nuploaded -- restart the Space to pick the packages up" | |
| def budget(ref, cond, keep, frames, apath, seed): | |
| """Seconds of GPU to ask ZeroGPU for. Three things ride on it, none of them the bill: quota | |
| is charged on the *actual* elapsed window (a request that never touches the GPU bills zero), | |
| so padding costs nothing. It is an admission check -- a visitor whose remaining Hub quota is | |
| under 2x this number is rejected before anything runs, because `xlarge` doubles server-side | |
| -- a queue-priority hint (shorter first), and a soft ceiling: when the truly-cold worker | |
| blew ~90 s past the old 135 s request, the render finished and the log got a warning, not a | |
| kill. Do not lean on that leniency; keep it an over-estimate. | |
| The slope is measured on one B200: 26 s at 124 frames on 480x832, 72 s at 243 on 480x864 -- | |
| ~0.39 s per extra frame, and that 480x864 slightly exceeds LADDER's largest area, so the | |
| slope over-covers every canvas. SLOWDOWN is this card's measured ratio to that B200 (42/26 s | |
| at 124 frames). H2D covers the 73.5 GiB host-to-device copy, billed inside the window but | |
| invisible to `render`'s timer because it runs in worker init: measured 22 s from warm page | |
| cache, 75 s typical, 154 s cold -- and one 261 s outlier (2026-09-05) that overran the | |
| whole request by ~40 s and, again, drew only the warning. | |
| The constants are set high ON PURPOSE. The admission check is also the | |
| visitor filter: at these values the shortest render asks for 247 s -- 494 s at admission -- | |
| which is more than a free ZeroGPU day and small change against a PRO day (2400 s). That, | |
| not a sign-in, is what keeps anonymous and free-tier GPU use out; lower the H2D/SLOWDOWN | |
| Space variables and the free tier comes back. | |
| """ | |
| if frames is None: | |
| # gr.State does not survive a Space restart; a page opened before one sends None for | |
| # every State. This runs at admission, before `render`, so it is the only place the | |
| # visitor can get a message instead of a bare TypeError (production, 2026-09-05). | |
| raise gr.Error("this session predates a Space restart -- run Prepare again") | |
| return 35 + H2D + round(SLOWDOWN * (26 + 0.39 * (frames - MIN_FRAMES))) | |
| def render(ref, cond, keep, frames, apath, seed): | |
| global ATTN | |
| # The real 73.5 GiB pageable copy on a cold worker. A reused warm worker pays nothing -- | |
| # but do not count on reuse: it requires the scheduler to return `idle`, a server-side | |
| # decision, and a worker that raised is dropped from the cache rather than re-registered. | |
| # `budget`'s H2D term assumes the pessimistic case. | |
| t_h2d = time.time() | |
| PIPE.to("cuda") | |
| h2d = time.time() - t_h2d | |
| # Workers are reused, so without this the reported peak is the high-water mark of every | |
| # render that worker has ever served, not this one's. | |
| torch.cuda.reset_peak_memory_stats() | |
| # cuDNN's attention kernel is a measured 2.05x on the denoise loop and does not move the | |
| # sample -- same weights, same seed, same take -- so it can be switched on without | |
| # re-rendering the bundled examples. It is selected here and not at import because there is | |
| # no GPU at import under ZeroGPU. `set_attention_backend` validates nothing for this backend | |
| # (`_check_attention_backend_requirements` has no branch for it), so an unserviceable shape | |
| # raises at the first attention call, which is why the pipeline call below can retry. That | |
| # 2.05x was measured on sm_100 and this card is sm_120. `ATTN=""` in the Space variables | |
| # backs it out without a code change. | |
| if ATTN: | |
| PIPE.transformer_ref.set_attention_backend(ATTN) | |
| wire_aoti() | |
| # `frames` arrives from `prepare` through State -- it cut `cond` to exactly that count, so | |
| # re-counting here would be a full ffprobe decode inside the billed window for a number the | |
| # parent already knows. | |
| frames = int(frames) | |
| w, h = probe(cond) | |
| PIPE.register_to_config(canvas_short_edge=min(w, h), | |
| canvas_max_pixels=min(w, h) * max(w, h), | |
| reference_image_short_edge=min(w, h)) | |
| # The reference order is frozen -- `<Video 1>` then `<Picture 1>` -- because that order | |
| # advances the shared rotary clock. The clip's soundtrack is still cleared off the reference, | |
| # as in training -- it is never conditioning -- but it is no longer discarded: when there is | |
| # one it goes into the *target* audio rows instead, pinned clean, below. | |
| # | |
| # An older comment here claimed "the audio you hear back is generated". That was wrong. This | |
| # finetune's audio head emits digital silence (RMS ~1e-5), which is why every render has | |
| # needed a manual remux; nothing is lost by overwriting those rows. | |
| video = MiniMaxH3VideoReference.from_file(cond) | |
| video.audio, video.sample_rate = None, None | |
| # Encode the kept soundtrack into the target audio rows. `PIPE.audio_vae` is already on the | |
| # card, and bf16 is free here: measured against an fp32 CPU encode the rows differ by 5.4% | |
| # relative L2, where speech and silence differ by 127% -- and the decoded waveforms are | |
| # indistinguishable (|amp| r 0.9188 vs 0.9191, the VAE round-trip dominating both). A second | |
| # fp32 copy would cost 578 MB of host RAM to buy nothing. | |
| # | |
| # `attention_backend()` rather than trusting statement order: the audio *encoder* dispatches | |
| # `is_causal=True`, which cuDNN has no kernel for, and `set_attention_backend` above is | |
| # process-wide despite looking per-model (its last line sets a classvar on | |
| # `_AttentionBackendRegistry`). Without this the encode dies in `_native_cudnn_attention` with | |
| # "No available kernel". The manager restores in `finally`, so denoise keeps the 2.05x kernel. | |
| # It works here only because the audio VAE's own `_attention_backend` is still None -- setting | |
| # a backend on the VAE itself would pass `backend=` to `dispatch_attention_fn` and bypass this. | |
| global PIN | |
| PIN = apath is not None | |
| audio_latents = None | |
| if PIN: | |
| wav = torch.frombuffer(bytearray(Path(apath).read_bytes()), | |
| dtype=torch.float32).view(-1, 2).T | |
| n_samp = round((frames + 1) / FPS * ASR) # one frame over, so the encoder has enough | |
| 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 | |
| audio_latents = rows.permute(0, 2, 1).contiguous() # (2, 32, n_lat) | |
| print(f"[gpu] pinned {wav.shape[1]/ASR:.2f}s of audio -> {tuple(audio_latents.shape)}", | |
| flush=True) | |
| kwargs = dict( | |
| prompt=FIXED["presentation"], | |
| references=[video, MiniMaxH3ImageReference.from_file(ref)], | |
| num_frames=frames, height=h, width=w, num_inference_steps=4, | |
| audio_latents=audio_latents, | |
| generator=torch.Generator().manual_seed(int(42 if seed is None else seed)), | |
| output=["videos", "audio", "sampling_rate"], | |
| ) | |
| t0 = time.time() | |
| try: | |
| result = PIPE(**kwargs) | |
| except Exception as e: # noqa: BLE001 -- one retry, only for the backend | |
| if not ATTN: | |
| raise | |
| # cuDNN rejects a shape at the *first* attention call, seconds into a budget of minutes, | |
| # so there is room to start over on the backend that is always correct. Latching ATTN | |
| # off means this worker retries at most once, and every later render skips cuDNN outright. | |
| print(f"[gpu] {ATTN} failed ({type(e).__name__}: {e}); retrying on the default backend", | |
| flush=True) | |
| ATTN = "" | |
| PIPE.transformer_ref.set_attention_backend("native") | |
| # The compiled blocks baked cuDNN at build time and ignore the line above -- eager is | |
| # the only backend change that can actually happen, so the retry drops them too. | |
| unwire_aoti() | |
| result = PIPE(**kwargs) | |
| took = time.time() - t0 | |
| work = Path(tempfile.mkdtemp()) | |
| raw = str(work / "raw.mp4") | |
| encode_video(result["videos"][0], fps=FPS, output_path=raw, | |
| audio=result["audio"][0], audio_sample_rate=result["sampling_rate"]) | |
| msg = (f"{took:.0f}s for 4 steps at {frames} frames ({keep:.1f}s out), peak " | |
| f"{torch.cuda.max_memory_allocated()/2**30:.1f} GiB") | |
| # Calibrate SLOWDOWN against `took` and H2D against `h2d`. | |
| print(f"[gpu] {msg}, weights onto card in {h2d:.0f}s " | |
| f"(budget assumed {budget(ref, cond, keep, frames, apath, seed)}s)", flush=True) | |
| # `keep` rides along with the path instead of being read back from State in `finish`: State | |
| # is read when each chain link *starts*, so a second Generate click during this render would | |
| # overwrite keep_s and the finished clip would be trimmed to the wrong length. | |
| return (raw, keep, apath), msg | |
| def finish(job): | |
| """Trim and watermark, on the parent's CPU. This used to run inside `render`, where every | |
| second is billed GPU time; x264 and drawtext need no card. `-t keep` drops the frozen tail | |
| the 17k+5 grid forced us to render, so the clip that comes back is exactly as long as the | |
| one that went in. drawtext needs an explicit fontfile: the container has no fontconfig | |
| default, and without one ffmpeg fails here, after the render -- the worst place to fail | |
| (the boot-time assert on FONT exists to keep it unreachable). | |
| """ | |
| raw, keep, apath = job | |
| out = str(Path(raw).parent / "viggle-animate.mp4") | |
| # With a pinned track the render's own audio rows do decode back to the upload's soundtrack, | |
| # but only through the audio VAE, at |amp| r 0.92. The original is on disk and exact, so it is | |
| # remuxed in instead. The timelines agree by construction: `prepare`'s fps filter retimes | |
| # without time-stretching, and it cut the audio to the same `keep` the video is trimmed to here. | |
| src = ["-f", "f32le", "-ar", str(ASR), "-ac", "2", "-i", apath] if apath else [] | |
| amap = ["-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac"] if apath else ["-c:a", "copy"] | |
| run("ffmpeg", "-y", "-loglevel", "error", "-i", raw, *src, "-t", f"{keep:.3f}", "-vf", | |
| f"drawtext=fontfile={FONT}:text='viggle-animate':x=w-tw-14:y=h-th-14:fontsize=22:" | |
| "fontcolor=white@0.85:shadowcolor=black@0.6:shadowx=1:shadowy=1", | |
| *amap, "-crf", "20", "-pix_fmt", "yuv420p", out) | |
| return out | |
| with gr.Blocks(title="Viggle-Animate") as demo: | |
| gr.Markdown( | |
| "# Viggle-Animate\n" | |
| "Put the person from a still into a driving video. No pose estimation, no " | |
| "segmentation, no masks — the only preprocessing is repainting the clip's first " | |
| f"frame. 4 sampling steps. Takes {MIN_UPLOAD_S:.0f}–{MAX_FRAMES/FPS:.0f} s of driving " | |
| f"video; anything longer is cut to {MAX_FRAMES/FPS:.1f} s. If the clip carries a " | |
| "soundtrack it is pinned into the render and the mouth tracks it; upload a silent clip " | |
| "and the mouth moves without saying anything.\n\n" | |
| f"[Model card and weights]({MODEL_CARD}) · [viggle.ai]({SITE})" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| v = gr.Video(label="driving video — motion, framing and background are kept; " | |
| "its soundtrack, if it has one, is what the mouth follows") | |
| c = gr.Image(label="character — who to put in it", type="filepath") | |
| p = gr.Image(label="…or frame 0, already painted (skips gpt-image-2)", | |
| type="filepath") | |
| go = gr.Button("Generate", variant="primary") | |
| seed = gr.Number(value=42, label="seed", precision=0) | |
| gr.Markdown(f"<sub>No sign-in. The first-frame repaint is paid by this Space — " | |
| f"{DAILY} a day per visitor. The render runs on your own ZeroGPU " | |
| "quota and asks for more than a free day holds, so in practice PRO " | |
| "and up.</sub>") | |
| with gr.Column(): | |
| shown = gr.Image(label="step 1: repainted first frame", type="filepath") | |
| out = gr.Video(label="step 2: render") | |
| log = gr.Textbox(label="", lines=2) | |
| # The examples come with frame 0 already painted, so clicking one exercises the model alone: | |
| # no OpenAI call, and the same seed gives the same render every time. `out` is listed as an | |
| # input on purpose -- that is what lets a click drop the finished render into the output slot | |
| # next to the two things that produced it, so the pair can be read without spending quota. | |
| gr.Examples( | |
| examples=[[str(EX / f"{d}_drive.mp4"), None, str(EX / f"{n}_ref.png"), | |
| str(EX / f"{n}_out.mp4")] for n, (d, _) in CASES.items()], | |
| inputs=[v, c, p, out], examples_per_page=len(CASES), | |
| label="examples — click one for its driving clip, its painted first frame and the render", | |
| ) | |
| # Rendering asks ZeroGPU for more quota than a free account's day holds (see `budget`), so | |
| # for most people arriving here this row is the only output they will see -- the page's | |
| # actual evidence rather than a footnote. | |
| gr.Markdown("### What it produces\n" | |
| "Every example above, rendered at 4 steps with seed 42. Each is a single paint " | |
| "and a single render; the first and the last three swap in a person. Only the " | |
| "first one has audio — it is the only one whose driving clip had any.") | |
| items = list(CASES.items()) | |
| # Rows are cut on aspect, not on a fixed width: the landscape lip-sync take on its own, then | |
| # six portrait swaps, the four landscape two-character takes, and the three human cases. | |
| # Mixing the two shapes in one row squeezes the portrait ones to a third of their height. | |
| for row in (items[0:1], items[1:4], items[4:7], items[7:11], items[11:14]): | |
| with gr.Row(): | |
| for n, (_, caption) in row: | |
| gr.Video(str(EX / f"{n}_out.mp4"), label=caption, interactive=False) | |
| # Maintenance, visible only when the AOTI_BUILD Space variable is set (and refused for | |
| # anyone but the owner even then): compiles the block packages on this hardware and uploads | |
| # them to the weights repo. Run once per torch/driver generation, then unset the variable. | |
| # The LoginButton is the only OAuth left in the app and exists purely for this gate: | |
| # `build_ui`'s `gr.OAuthProfile` stays None without a sign-in, and stays None for everyone | |
| # unless `hf_oauth: true` is in the README front matter -- keep both or the button refuses | |
| # even the owner. | |
| if os.environ.get("AOTI_BUILD") == "1": | |
| with gr.Accordion("maintenance -- build AOTI packages", open=False): | |
| gr.LoginButton(size="sm") | |
| b_go = gr.Button("compile the repeated blocks on this GPU (<=18 min of " | |
| "xlarge quota), upload to the weights repo") | |
| b_log = gr.Textbox(label="build log", lines=12) | |
| b_go.click(build_ui, [], [b_log]) | |
| # Repeated at the bottom because the examples and the "What it produces" rows push the | |
| # header well off screen -- this is where a visitor actually finishes reading. | |
| gr.Markdown(f"<sub>[Model card and weights]({MODEL_CARD}) · [viggle.ai]({SITE})</sub>") | |
| ref_s, cond_s, keep_s, frames_s, audio_s, raw_s = (gr.State(), gr.State(), gr.State(), | |
| gr.State(), gr.State(), gr.State()) | |
| # `.success`, not `.then`: `prepare` raises gr.Error for a missing upload, a too-short clip | |
| # or a spent daily quota, and `.then` would run the render anyway -- on stale State from the | |
| # previous click, or on None, which is a several-minute GPU job started by a failed one. | |
| # `finish` runs in the parent, but the raw path crosses fine: worker and parent share /tmp. | |
| # `finish` runs outside the queue: it is seconds of CPU ffmpeg, and inside the queue a | |
| # "Queue is full" rejection would land *after* the GPU time is already spent. | |
| (go.click(prepare, [v, c, p], | |
| [ref_s, cond_s, keep_s, frames_s, audio_s, shown, out, log]) | |
| .success(render, [ref_s, cond_s, keep_s, frames_s, audio_s, seed], [raw_s, log]) | |
| .success(finish, [raw_s], [out], queue=False)) | |
| # The bundled media sits next to app.py rather than in Gradio's temp dir, so it has to be | |
| # named explicitly or the file route refuses to serve it. | |
| # `max_file_size` matters more than it looks: gradio writes the upload to disk before | |
| # `prepare` can refuse it, and that disk also holds 80 GB of weights. The default is | |
| # None, i.e. unlimited. 10 s of 4K is comfortably under 200 MB. | |
| # Guarded so ../space_viggle_wf/workflow_app.py can import this module (the pipeline and the | |
| # boot-time weight load) without starting a second server. | |
| if __name__ == "__main__": | |
| demo.queue(max_size=32).launch(show_error=True, allowed_paths=[str(EX)], | |
| max_file_size="200mb") | |