Download inference/sample.py from Viggle/Viggle-Remake: direct link, hf CLI and curl.
- Browser
- Download file 17.7 kB
-
https://huggingface.co/Viggle/Viggle-Remake/resolve/main/inference/sample.py
- Command line
-
hf download hf://Viggle/Viggle-Remake/inference/sample.py
-
curl -L -o sample.py https://huggingface.co/Viggle/Viggle-Remake/resolve/main/inference/sample.py
17.7 kB
| # SPDX-License-Identifier: Apache-2.0 | |
| """Viggle Remake: re-render a video with the people (swap) or the scene (edit) changed. | |
| python inference/sample.py --mode swap --video source.mp4 --prompt prompt.txt \ | |
| --ref ref1.png [--ref ref2.png ...] [--pin 0 [--pin 96 ...]] \ | |
| --frames 124 --short-edge 480 --seed 0 \ | |
| --base MiniMax-H3 --lora-repo Viggle-Remake --out out.mp4 | |
| `--video` supplies motion, camera, framing and everything the prompt does not change. Each | |
| `--ref` is an image the prompt refers to as <Picture N>. `--pin` is optional and, when given, | |
| appears once per `--ref`: it pins that ref as a keyframe at that real frame of the output for | |
| the first 90% of the denoise (`--keyframe-until`), which is how a painted frame of the source | |
| becomes an anchor rather than only a reference. | |
| The target audio rows hold the source soundtrack (`--audio`, defaults to `--video`) as a clean | |
| latent for the whole denoise, so mouths follow the real speech and the output carries that | |
| soundtrack back through the audio VAE. | |
| Memory order: the 64-layer Qwen3-VL conditioner (~59 GiB) is loaded, encodes the prompt once, | |
| and is freed before the pipeline loads, so the two never share the card. | |
| Runs on upstream diffusers d6726f3. The keyframe pin (`ref2va` references *plus* keyframes inside | |
| the generated span, the layout the LoRAs are run with) is not in | |
| upstream; the three small patches under "keyframes on a ref2va layout" below add it. | |
| """ | |
| import argparse | |
| import gc | |
| import subprocess | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| import av | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration | |
| 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.denoise import MiniMaxH3LoopSchedulerStep, MiniMaxH3Ref2VALoopDenoiser | |
| from diffusers.modular_pipelines.minimax_h3.encoders import (MiniMaxH3Ref2VAReferenceEncoderStep, | |
| MiniMaxH3Ref2VATextEncoderStep, encode_vae_condition, | |
| get_qwen3vl_prompt_embeds) | |
| from diffusers.modular_pipelines.minimax_h3.modular_pipeline import audio_latent_num_frames | |
| from diffusers.utils.export_utils import encode_video | |
| # 6-step grids, in the video stream's sigma; each stream's own shift is applied below | |
| SIGMAS = {"swap": [1, .990476, .965243, .923077, .8, .631579, 0], | |
| "edit": [1, .983607, .96, .923077, .857143, .705882, 0]} | |
| TEXT_TAG = 1 | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--mode", required=True, choices=["swap", "edit"]) | |
| parser.add_argument("--video", required=True, help="the source clip") | |
| parser.add_argument("--audio", default=None, help="the soundtrack pinned into the output; defaults to --video") | |
| parser.add_argument("--prompt", required=True, help="text file with the six-section prompt, read verbatim") | |
| parser.add_argument("--ref", action="append", required=True, help="reference image, <Picture N> in the prompt") | |
| parser.add_argument("--pin", action="append", type=int, default=[], | |
| help="output frame each --ref is pinned to as a keyframe; absent, or once per --ref") | |
| parser.add_argument("--frames", type=int, default=124, help="17n+5 in 90..345, at 24 fps") | |
| parser.add_argument("--short-edge", type=int, default=480, help="480 (832x480 for 16:9) or 736 (1280x736)") | |
| parser.add_argument("--seed", type=int, default=0) | |
| parser.add_argument("--base", required=True, help="local copy of MiniMaxAI/MiniMax-H3") | |
| parser.add_argument("--lora-repo", required=True, help="local copy of Viggle/Viggle-Remake") | |
| parser.add_argument("--rank", type=int, default=256, choices=[256, 128, 64], | |
| help="256 loads lora/<mode>; 128 or 64 load the smaller lora_lowrank/<mode>-r<rank>") | |
| parser.add_argument("--keyframe-until", type=float, default=0.9, help="drop the pinned keyframes once 1-t reaches this") | |
| parser.add_argument("--attn-backend", default="_native_cudnn") | |
| parser.add_argument("--out", required=True) | |
| args = parser.parse_args() | |
| nf, se, refs, pin = args.frames, args.short_edge, args.ref, args.pin | |
| assert 90 <= nf <= 345 and (nf - 5) % 17 == 0, f"--frames must be 17n+5 in 90..345, got {nf}" | |
| assert se % 32 == 0, f"--short-edge must be a multiple of 32, got {se}" | |
| assert len(pin) in (0, len(refs)), "give --pin once per --ref, or not at all" | |
| assert all(0 <= k < nf for k in pin), f"--pin must be in [0, {nf})" | |
| prompt = Path(args.prompt).read_text() | |
| sigmas = SIGMAS[args.mode] | |
| # The source on its canvas. A clip already at the canvas (short edge == --short-edge, both sides multiples of 32) is used | |
| # as is; anything else is cover-cropped onto the 16:9-class canvas the way our renderer does (832x480, 1280x736). | |
| with av.open(args.video) as c: | |
| w, h = c.streams.video[0].codec_context.width, c.streams.video[0].codec_context.height | |
| video_path = args.video | |
| if min(w, h) != se or w % 32 or h % 32: | |
| long = se * 16 // 9 // 32 * 32 | |
| w, h = (long, se) if w >= h else (se, long) | |
| video_path = str(Path(tempfile.mkdtemp()) / "cond.mp4") | |
| subprocess.run(["ffmpeg", "-v", "error", "-i", args.video, "-vf", | |
| f"fps=24,scale={w}:{h}:force_original_aspect_ratio=increase,crop={w}:{h}", "-frames:v", str(nf), | |
| "-an", "-c:v", "libx264", "-crf", "12", "-pix_fmt", "yuv420p", video_path], check=True) | |
| video = MiniMaxH3VideoReference.from_file(video_path) | |
| video.audio, video.sample_rate = None, None # the soundtrack goes in as the target audio pin, not as reference audio | |
| height, width = video.frames.shape[1], video.frames.shape[2] | |
| # --- prompt -> embeds. The presentation is text only: a label per 2 fps block of the source, one per ref, then the | |
| # prompt, each piece tokenized on its own. H3 reads hidden_states[50] of the 64-layer conditioner. | |
| t0 = time.time() | |
| processor = AutoProcessor.from_pretrained(args.base, subfolder="processor") | |
| tokenizer = AutoTokenizer.from_pretrained(args.base, subfolder="tokenizer") | |
| text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(args.base, subfolder="text_encoder", | |
| dtype=torch.bfloat16).to("cuda") | |
| t1 = time.time() | |
| _, stamps = MiniMaxH3Ref2VATextEncoderStep._sample_video_condition_frames( | |
| np.zeros((nf, 1, 1, 3), dtype=np.uint8), 24.0, 2.0, processor.video_processor.temporal_patch_size) | |
| segments = ["<Video 1>: ", *[f"<{s:.1f} seconds>" for s in stamps], *[f"<Picture {i + 1}>: " for i in range(len(refs))], | |
| prompt] | |
| token_ids = [i for s in segments for i in tokenizer(s, add_special_tokens=False)["input_ids"]] | |
| with torch.no_grad(): | |
| prompt_embeds = get_qwen3vl_prompt_embeds(text_encoder, processor, token_ids, device=torch.device("cuda"), | |
| dtype=torch.bfloat16).cpu() | |
| text_token_tags = torch.full((len(token_ids),), TEXT_TAG, dtype=torch.long) | |
| t2 = time.time() | |
| print(f"text encoder: load {t1 - t0:.0f}s, encode {t2 - t1:.1f}s, {len(token_ids)} tokens, " | |
| f"peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB", flush=True) | |
| del text_encoder | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| # --- pipeline: the LoRA sits on the fl2va `transformer`, loaded into the ref2va slot | |
| pipe = ModularPipeline.from_pretrained(args.base, workflow="ref2va") | |
| type(pipe).min_duration = property(lambda self: 90 / 24) # the adapters run from 90 frames; the base floor is 5 s | |
| pipe.load_components(names=["vae", "audio_vae", "scheduler", "audio_scheduler"], pretrained_model_name_or_path=args.base, | |
| dtype=torch.bfloat16) | |
| pipe.transformer_ref = MiniMaxH3Transformer3DModel.from_pretrained(args.base, subfolder="transformer", | |
| torch_dtype=torch.bfloat16) | |
| # prefix=None: the saved keys are bare module names; the default "transformer." filter loads nothing and only warns | |
| lora = Path(args.lora_repo) / (f"lora/{args.mode}" if args.rank == 256 else f"lora_lowrank/{args.mode}-r{args.rank}") | |
| pipe.transformer_ref.load_lora_adapter(str(lora), | |
| weight_name="pytorch_lora_weights.safetensors", prefix=None) | |
| pipe.transformer_ref.set_attention_backend(args.attn_backend) # process-wide: the audio encode below opts out | |
| pipe.to("cuda") | |
| pipe.register_to_config(canvas_short_edge=se, canvas_max_pixels=se * max(height, width), reference_image_short_edge=se) | |
| sh = lambda u, s: s * u / (1 + (s - 1) * u) # noqa: E731 | |
| for sched, s in ((pipe.scheduler, 12.0), (pipe.audio_scheduler, 3.0)): | |
| grid = [sh(x / (12 - 11 * x), s) for x in sigmas] # video sigma -> u -> this stream's sigma | |
| sched.set_timesteps = (lambda f, g: lambda *a, **k: f(sigmas=g, device=k.get("device")))(sched.set_timesteps, grid) | |
| t3 = time.time() | |
| print(f"pipeline on cuda in {t3 - t2:.0f}s", flush=True) | |
| def use_fixed_embeds(self, components, state): | |
| block_state = self.get_block_state(state) | |
| block_state.prompt_embeds = prompt_embeds.to(components._execution_device, torch.bfloat16) | |
| block_state.text_token_tags = text_token_tags | |
| self.set_block_state(state, block_state) | |
| return components, state | |
| MiniMaxH3Ref2VATextEncoderStep.__call__ = use_fixed_embeds | |
| # audio pin: the target audio rows are the source soundtrack as a clean latent (positional arg 6 = audio_timestep, and | |
| # H3's t = 1 is clean), and only the video rows are stepped, so they stay clean | |
| _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. A keyframe block is packed after the references, right before the target audio, and | |
| # is pinned *inside* the target's rotary span at its real frame index rather than advancing the clock like a reference. | |
| KF = {"images": [], "anchors": list(pin)} | |
| _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) | |
| # --keyframe-until: once 1 - t reaches it the keyframe blocks are cut out of every per-row input, leaving the plain | |
| # ref2va sequence; their velocity comes back as zeros, which the scheduler never reads (conditioning rows) | |
| r, (_, ph, pw) = pipe.vae_spatial_compression_ratio, pipe.patch_size | |
| nkf = len(pin) * (height // r // ph) * (width // r // pw) | |
| _denoise = MiniMaxH3Ref2VALoopDenoiser.__call__ | |
| def kf_release(self, components, block_state, i, t): | |
| if not nkf or 1 - float(t) >= args.keyframe_until: | |
| return _denoise(self, components, block_state, i, t) | |
| ncv = 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 | |
| # keyframes onto the canvas the way our renderer does: the first one stretched, the rest cover-cropped (rounded | |
| # size, (src - dst) // 2 centring -- not VaeImageProcessor's "crop" arithmetic, which differs by a pixel on some ratios) | |
| if pin: | |
| for i, im in enumerate(Image.open(x).convert("RGB") for x in refs): | |
| if im.size != (width, height) and i == 0: | |
| im = pipe.image_processor.resize(im, height=height, width=width) | |
| elif im.size != (width, height): | |
| scale = max(width / im.size[0], height / im.size[1]) | |
| size = (max(width, round(im.size[0] * scale)), max(height, round(im.size[1] * scale))) | |
| left, top = max(0, (size[0] - width) // 2), max(0, (size[1] - height) // 2) | |
| im = im.resize(size, Image.Resampling.LANCZOS).crop((left, top, left + width, top + height)) | |
| KF["images"].append(im) | |
| # --- the soundtrack as clean target audio latents | |
| asr = pipe.audio_vae.config.sampling_rate | |
| pcm = subprocess.run(["ffmpeg", "-loglevel", "error", "-i", args.audio or args.video, "-t", f"{nf / 24:.3f}", "-vn", | |
| "-f", "f32le", "-ac", "2", "-ar", str(asr), "-"], capture_output=True, check=True).stdout | |
| wav = torch.frombuffer(bytearray(pcm), dtype=torch.float32).view(-1, 2).T | |
| n_samp = round((nf + 1) / 24 * asr) | |
| 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(nf)] - mean) / std | |
| audio_latents = rows.permute(0, 2, 1).contiguous() | |
| torch.cuda.reset_peak_memory_stats() | |
| t4 = time.time() | |
| result = pipe(prompt="".join(segments), references=[video, *[MiniMaxH3ImageReference.from_file(x) for x in refs]], | |
| num_frames=nf, height=height, width=width, num_inference_steps=len(sigmas) - 1, | |
| generator=torch.Generator().manual_seed(args.seed), audio_latents=audio_latents, | |
| output=["videos", "audio", "sampling_rate"]) | |
| t5 = time.time() | |
| Path(args.out).parent.mkdir(parents=True, exist_ok=True) | |
| encode_video(result["videos"][0], fps=24, output_path=args.out, audio=result["audio"][0], | |
| audio_sample_rate=result["sampling_rate"]) | |
| print(f"render {t5 - t4:.1f}s at {width}x{height}x{nf}, peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB " | |
| f"allocated / {torch.cuda.max_memory_reserved() / 2**30:.1f} GiB reserved -> {args.out}", flush=True) | |