Masked video+audio inpainting blocks for MiniMax-H3

README.md

@@ -0,0 +1,129 @@
+---
+license: apache-2.0
+base_model: MiniMaxAI/MiniMax-H3
+base_model_relation: adapter
+library_name: diffusers
+tags:
+  - modular-diffusers
+  - minimax-h3
+  - inpainting
+  - video-to-video
+  - audio-video
+pipeline_tag: video-to-video
+---
+
+# MiniMax-H3 — masked video and audio inpainting
+
+Modular Diffusers blocks that repaint part of a clip with MiniMax-H3 and keep the rest, including the soundtrack.
+**No dedicated checkpoint, no adapter, no extra input channel** — the mask becomes a per-row timestep, which the model
+already had.
+
+![source above, repainted below](assets/example.webp)
+
+<sub>Above: the plate. Below: the same clip with the animal replaced from one reference photo, in 6 steps. The forest,
+the snow, the camera push and the original soundtrack are untouched.</sub>
+
+## Use it
+
+```python
+import torch
+from diffusers.modular_pipelines import ModularPipelineBlocks
+from diffusers.modular_pipelines.minimax_h3.references import MiniMaxH3ImageReference
+
+blocks = ModularPipelineBlocks.from_pretrained(
+    "diffusers-modular/minimax-h3-inpainting", trust_remote_code=True
+)
+pipe = blocks.init_pipeline("MiniMaxAI/MiniMax-H3")
+pipe.load_components(dtype=torch.bfloat16)
+pipe.to("cuda")
+
+state = pipe(
+    prompt="<Picture 1> the man from the picture, walking through deep snow in a pine forest",
+    references=[MiniMaxH3ImageReference.from_file("subject.png")],
+    source_video=frames,          # (num_frames, height, width, 3) uint8
+    source_fps=24,
+    mask=mask,                    # (num_frames, height, width) — 1 repaints, 0 preserves
+    source_audio=waveform,        # optional; preserved whole unless `audio_mask` says otherwise
+    source_audio_sample_rate=48000,
+    num_inference_steps=28,
+    generator=torch.Generator("cpu").manual_seed(0),
+)
+video, audio = state.get("videos")[0], state.get("audio")[0]
+```
+
+`MiniMaxH3Ref2VAInpaintGeneratorBlocks` is the same thing without the text-encoder step, for split deployments where
+the encoder lives elsewhere and `prompt_embeds` / `text_token_tags` are the wire format.
+
+## How it works
+
+MiniMax-H3 denoises one packed sequence in which **every row carries its own timestep** — that is how a keyframe
+anchor sits at `t = 0.999`, essentially clean, beside target rows still stepping down the schedule. Nothing says which
+rows may do that, so pointing it at an arbitrary subset of the target rows *is* inpainting.
+
+| mask | row timestep | content |
+|---|---|---|
+| `1` — repaint | the schedule's `t` | the model's |
+| `0` — preserve | `max(t, 0.999)` video, `1.0` audio | the source, clean |
+| feathered | `1 − m·σ` | blended to that level |
+
+This matters because the usual recipe — re-noise the source to the current sigma and blend — is *off-distribution*
+here: it hands the model a target row claiming timestep `t` while holding content at a level it never saw paired with
+that label. Presenting preserved rows as conditioning is a distribution the checkpoint knows well.
+
+Verified: with the mask all ones the blocks reproduce stock `ref2va` **bit for bit**, video and audio. On real
+weights the preserved region comes back at 1.44/255 from the source — the autoencoder round-trip floor — against
+40.0/255 inside the mask, and the soundtrack at cosine 0.972.
+
+## The mask lands on three grids
+
+A generic resize reproduces none of them, and getting any one wrong is a silent quality bug:
+
+- **spatially** — the VAE's 16× compression, then the transformer's 2×2 patch. A row is one token: it carries one
+  timestep and is written back whole, so a 2×2 latent patch is the finest a mask can be.
+- **temporally** — the VAE's chunked causal grouping, `(1, 4, 4, 4, 4)` repeating every 17 frames. Not uniform.
+- **on the audio clock** — 40 latents per second, *not* 24 frames per second. Aligning an audio mask to the video
+  grid is what puts a masked soundtrack out of sync.
+
+`pixel_mask_to_row_mask` and `audio_mask_to_row_mask` do this; every reduction is a maximum, so a row regenerates as
+much as the most-masked pixel it covers asks it to.
+
+## Use hard masks
+
+![plate, feathered mask, hard mask](assets/edge.webp)
+
+<sub>Plate · feathered mask · hard mask, at the same boundary.</sub>
+
+A feathered mask leaves its edge rows at intermediate timesteps holding a *mixture* of source and repaint — lower
+contrast than either. Paste that through an upscale and crossfade it into a sharp plate and you get a visible band
+along the mask, as in the middle panel. Squaring the mask off and generating at the plate's own size removes it. The
+paste's own feather is what should hide the join.
+
+## Give the mask room
+
+![a tracked segmentation, grown](assets/mask.webp)
+
+Mask geometry decides what a prompt can do. A box fitted to a walking quadruped is a quadruped-shaped hole: asked for
+a person, the model will put one in it *on all fours* rather than contradict the border it was told to preserve. Only
+a mask with a standing footprint lets it stand up. When you are replacing a subject rather than editing one, grow the
+mask well past the outline.
+
+`crop.py` ships the other half of the practical workflow: one stable box around everything the mask ever touches,
+a canvas that never upscales it, and a feathered paste back into the plate. Cost is set by the canvas, not by how
+much of the frame changes, so cropping to the subject is the memory lever.
+
+## Practice notes
+
+- **Keep the soundtrack** and the model animates to the words already there. That is the lip-sync recipe.
+- **Per-shot prompting is unavoidable.** Masking makes the prompt less strict, not optional.
+- **`ref2va` needs at least one reference.** Prompt-only object removal is the `t2va` partition's job.
+- **5–15 s per pass.** Longer clips have to be inpainted in segments.
+- **The decoder is not perfectly local.** Preserved latents are preserved exactly, but the video decoder is a
+  36-layer attention stack, so a change inside the mask moves decoded pixels just outside it — 5.5/255 within 8 px,
+  1.2/255 by 32 px, gone by 128 px. Confining the paste to the mask discards that halo.
+
+## Credits
+
+The workflow these blocks reproduce was worked out by the ComfyUI community — Ablejones, Nekodificador, and drozbay,
+whose [MaskVidExperiments](https://github.com/drozbay/MaskVidExperiments) covers the same ground. **No code is carried
+over from it** (it is GPLv3); the mask geometry here is re-derived from the checkpoint's own constants. ComfyUI's own
+MiniMax-H3 support arrives at the same "the mask is a timestep" formulation independently.

crop.py

@@ -0,0 +1,222 @@
+"""Crop a masked subject out of a clip, inpaint it small, paste it back.
+
+MiniMax-H3 generates on a 768-short-edge canvas and attends over the whole packed sequence at once, so the cost of a
+request is set by its canvas rather than by how much of the frame actually changes. Repainting a face in a 4K plate at
+full resolution is mostly spent on pixels the mask preserves — which is why every workflow in the wild crops to the
+subject first, and why "reduce the resolution until it stops running out of memory" is the usual advice.
+
+This is the static half of that: one box around everything the mask ever touches, held for the whole clip. A box that
+never moves cannot be read as camera motion, which is the failure mode of cropping each frame to its own subject; the
+cost is that a subject crossing the frame drags the box out to cover its whole travel. Tracking the subject with a box
+that moves as little as possible does better on those clips and is a much larger piece of work.
+
+Nothing here touches the model. Crop before the pipeline, paste after.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from PIL import Image, ImageFilter
+
+
+def mask_bounding_box(
+    mask: np.ndarray,
+    frame_height: int,
+    frame_width: int,
+    crop_scale: float = 0.5,
+    multiple: int = 32,
+    threshold: float = 0.02,
+    min_aspect_ratio: float = 0.25,
+    max_aspect_ratio: float = 4.0,
+) -> tuple[int, int, int, int]:
+    r"""
+    One box around everything the mask touches anywhere in the clip.
+
+    Args:
+        mask (`np.ndarray` of shape `(num_frames, height, width)`): The mask, over `[0, 1]`.
+        frame_height (`int`), frame_width (`int`): The frame the box is cut from, which it is clamped to.
+        crop_scale (`float`, defaults to 0.5):
+            Padding around the subject, as a fraction of its size. The model needs context around what it repaints —
+            a box cut tight to a head gives it nothing to match lighting or motion against.
+        multiple (`int`, defaults to 32): What both axes are rounded up to, i.e. the pipeline's `canvas_multiple`.
+        threshold (`float`, defaults to 0.02): Mask values at or below this are treated as unmasked.
+        min_aspect_ratio (`float`, defaults to 0.25), max_aspect_ratio (`float`, defaults to 4.0):
+            The ratios MiniMax-H3 was trained over. A box outside them is widened on its short axis.
+
+    Returns:
+        `tuple[int, int, int, int]`: the box as `(top, left, height, width)`, in source pixels.
+    """
+    if mask.ndim != 3:
+        raise ValueError(f"A mask must be `(num_frames, height, width)`, got {tuple(mask.shape)}.")
+
+    covered = mask.max(axis=0) > threshold
+    if not covered.any():
+        raise ValueError("The mask is empty: there is nothing to inpaint.")
+    rows = np.flatnonzero(covered.any(axis=1))
+    cols = np.flatnonzero(covered.any(axis=0))
+    top, bottom = int(rows[0]), int(rows[-1]) + 1
+    left, right = int(cols[0]), int(cols[-1]) + 1
+
+    # Padding, as a share of the subject rather than a fixed margin, so it scales with the shot.
+    pad_y = (bottom - top) * crop_scale / 2.0
+    pad_x = (right - left) * crop_scale / 2.0
+    top, bottom = top - pad_y, bottom + pad_y
+    left, right = left - pad_x, right + pad_x
+
+    height, width = bottom - top, right - left
+    # Back inside the ratios the checkpoint was trained over, by growing the short axis rather than cutting the long
+    # one — the subject stays fully inside the box either way.
+    ratio = width / height
+    if ratio < min_aspect_ratio:
+        width = height * min_aspect_ratio
+    elif ratio > max_aspect_ratio:
+        height = width / max_aspect_ratio
+
+    return _snap_box(
+        (top + bottom) / 2.0, (left + right) / 2.0, height, width, frame_height, frame_width, multiple
+    )
+
+
+def _snap_box(
+    center_y: float,
+    center_x: float,
+    height: float,
+    width: float,
+    frame_height: int,
+    frame_width: int,
+    multiple: int,
+) -> tuple[int, int, int, int]:
+    """A centred box onto the `multiple` grid and inside the frame, keeping the subject centred where it can."""
+    height = min(int(np.ceil(height / multiple)) * multiple, (frame_height // multiple) * multiple)
+    width = min(int(np.ceil(width / multiple)) * multiple, (frame_width // multiple) * multiple)
+    height, width = max(height, multiple), max(width, multiple)
+
+    top = int(round(center_y - height / 2.0))
+    left = int(round(center_x - width / 2.0))
+    # A box pushed off the edge slides back in rather than being cut down: its size is already on the grid, and
+    # shrinking it here would drop part of the subject.
+    top = max(0, min(top, frame_height - height))
+    left = max(0, min(left, frame_width - width))
+    return top, left, height, width
+
+
+def canvas_for_box(
+    box_height: int,
+    box_width: int,
+    multiple: int = 32,
+    short_edge: int = 768,
+    max_pixels: int = 768 * 1344,
+) -> tuple[int, int]:
+    r"""
+    The canvas to generate a box at.
+
+    Same rule the pipeline applies to any request — short edge first, area capped, both axes rounded — with one
+    addition: a box smaller than the canvas is generated at its own size rather than upscaled. Repainting a 320-pixel
+    head at 768 and scaling it back down spends the difference on nothing.
+
+    Lower `short_edge` and `max_pixels` together to trade quality for memory; the community workflows run at roughly
+    0.4-0.5 MP on 24 GB cards.
+
+    Args:
+        box_height (`int`), box_width (`int`): The box, as [`mask_bounding_box`] returned it.
+        multiple (`int`, defaults to 32): What both axes round to.
+        short_edge (`int`, defaults to 768), max_pixels (`int`, defaults to `768 * 1344`): The canvas budget.
+
+    Returns:
+        `tuple[int, int]`: the `(height, width)` to generate at.
+    """
+    scale = short_edge / min(box_height, box_width)
+    if box_height * box_width * scale * scale > max_pixels:
+        scale = (max_pixels / (box_height * box_width)) ** 0.5
+    scale = min(scale, 1.0)
+
+    height = max(multiple, round(box_height * scale / multiple) * multiple)
+    width = max(multiple, round(box_width * scale / multiple) * multiple)
+    return height, width
+
+
+def crop(array: np.ndarray, box: tuple[int, int, int, int]) -> np.ndarray:
+    r"""Cut `box` out of a `(num_frames, height, width, ...)` clip or mask. Pixel-exact, no resampling."""
+    top, left, height, width = box
+    return array[:, top : top + height, left : left + width]
+
+
+def paste_back(
+    frames: np.ndarray,
+    generated: np.ndarray,
+    box: tuple[int, int, int, int],
+    mask: np.ndarray | None = None,
+    feather: int = 8,
+) -> np.ndarray:
+    r"""
+    Paste an inpainted crop back into the frames it came from.
+
+    Args:
+        frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The original `uint8` clip.
+        generated (`np.ndarray` of shape `(num_frames, box_height, box_width, 3)` or the canvas size):
+            The inpainted crop, rescaled to the box if it was generated at another size.
+        box (`tuple[int, int, int, int]`): The box, as `(top, left, height, width)`.
+        mask (`np.ndarray` of shape `(num_frames, box_height, box_width)`, *optional*):
+            Confine the paste to the mask, blurred by `feather`. Without it the whole box is pasted, feathered only
+            at its border. Confining is the safer default on a static plate: the video VAE's decode is not local, so
+            a repainted region moves its surroundings slightly even where the mask preserved the latents exactly.
+        feather (`int`, defaults to 8): Width of the blend ramp in box pixels.
+
+    Returns:
+        `np.ndarray` of shape `(num_frames, height, width, 3)`: the `uint8` clip with the crop pasted in.
+    """
+    top, left, box_height, box_width = box
+    num_frames = frames.shape[0]
+    if generated.shape[0] != num_frames:
+        raise ValueError(f"The clip has {num_frames} frames and the inpainted crop {generated.shape[0]}.")
+
+    if generated.shape[1:3] != (box_height, box_width):
+        generated = np.stack(
+            [
+                np.asarray(Image.fromarray(frame).resize((box_width, box_height), Image.Resampling.LANCZOS))
+                for frame in generated
+            ]
+        )
+
+    weight = _border_ramp(box_height, box_width, feather)[None]
+    if mask is not None:
+        if mask.shape[1:3] != (box_height, box_width):
+            raise ValueError(
+                f"The mask is {mask.shape[1:3]} and the box {(box_height, box_width)}; crop the mask with the clip."
+            )
+        weight = weight * _feather_mask(mask, feather)
+
+    out = frames.copy()
+    region = out[:, top : top + box_height, left : left + box_width].astype(np.float32)
+    blended = region + weight[..., None] * (generated.astype(np.float32) - region)
+    out[:, top : top + box_height, left : left + box_width] = blended.round().clip(0, 255).astype(np.uint8)
+    return out
+
+
+def _border_ramp(height: int, width: int, feather: int) -> np.ndarray:
+    """1 inside the box, ramping to 0 over `feather` pixels at its edges."""
+    if feather <= 0:
+        return np.ones((height, width), dtype=np.float32)
+
+    def axis(length):
+        edge = np.minimum(np.arange(length), np.arange(length)[::-1]).astype(np.float32)
+        return np.clip((edge + 0.5) / feather, 0.0, 1.0)
+
+    return axis(height)[:, None] * axis(width)[None, :]
+
+
+def _feather_mask(mask: np.ndarray, feather: int) -> np.ndarray:
+    """The mask, blurred outwards so the paste fades into the plate rather than showing its own edge."""
+    if feather <= 0:
+        return mask.astype(np.float32)
+    blurred = []
+    for frame in mask:
+        image = Image.fromarray((frame * 255.0).clip(0, 255).astype(np.uint8))
+        image = image.filter(ImageFilter.MaxFilter(_odd(feather))).filter(ImageFilter.GaussianBlur(feather / 2.0))
+        blurred.append(np.asarray(image, dtype=np.float32) / 255.0)
+    return np.stack(blurred)
+
+
+def _odd(value: int) -> int:
+    """`MaxFilter` takes an odd kernel size, and 1 is a no-op rather than an error."""
+    return max(1, int(value) | 1)

minimax_h3_inpaint_blocks.py

@@ -0,0 +1,1046 @@
+"""Masked video+audio inpainting for MiniMax-H3, on the `ref2va` reference workflow.
+
+`ref2va` conditions on reference images, clips and soundtracks and then generates a whole new video. This module keeps
+that conditioning and makes the *target* partly given: a source clip is encoded, a mask says which of its rows the
+model repaints, and everything the mask preserves is presented to the transformer the way a keyframe anchor is —
+clean content pinned at the visual-conditioning timestep `0.999`, rather than the source noised down to the step's own
+sigma.
+
+That distinction is the whole trick, and it is what ComfyUI's MiniMax-H3 masking does too (`scale_latent_inpaint`
+returns the anchor rather than the schedule's noise level, and the DiT relabels masked rows). MiniMax-H3 was trained
+with clean visual conditioning at `0.999` and has never seen a partially denoised target row claiming to be something
+else, so ordinary "re-noise the original to sigma_t" inpainting is off-distribution for it.
+
+Three things follow, and they are the three blocks below:
+
+* **the mask lands on rows, not pixels.** A row is one `2 x 2` latent patch of one latent frame, and the latent
+  frames are grouped from pixel frames on the VAE's `(1, 4, 4, 4, 4)` cycle. `h3_inpaint_masks` does that reduction;
+  a generic resize does not reproduce it, and on the audio side the clock is 40 latents/s rather than 24 fps.
+* **preserved rows are relabelled.** MiniMax-H3 already carries a per-row timestep vector — that is how a keyframe
+  anchor rides at `0.999` beside target rows stepping down the schedule — so a mask value `m` simply places its row at
+  `1 - m * sigma`, clamped at the conditioning timestep. A fully preserved row lands exactly on `0.999`; a fully
+  generated one on the schedule; a feathered one in between, honestly labelled.
+* **preserved rows are written back every step.** The scheduler still steps them, and the write-back overrules it, the
+  way `MiniMaxH3AudioConditionStep` imposes a given soundtrack in the audio-to-video Space.
+
+The audio stream is masked on the same footing. Preserving the whole soundtrack — the default when a source
+soundtrack is passed — is what makes the model animate a mouth *to the words that are already there* rather than
+inventing new ones, which is the "keep the original spoken text" recipe from the community workflows.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from diffusers import attention_backend
+from diffusers.models import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    MiniMaxH3PrepareConditionLatentsStep,
+    MiniMaxH3PrepareLatentsStep,
+    MiniMaxH3Ref2VAPrepareLatentsStep,
+    MiniMaxH3Ref2VAPrepareLayoutStep,
+    MiniMaxH3SetTimestepsStep,
+    patchify_video_latents,
+)
+from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
+from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
+from diffusers.modular_pipelines.minimax_h3.denoise import (
+    MiniMaxH3DenoiseLoopWrapper,
+    MiniMaxH3LoopSchedulerStep,
+    MiniMaxH3Ref2VALoopDenoiser,
+)
+from diffusers.modular_pipelines.minimax_h3.encoders import (
+    MiniMaxH3Ref2VAReferenceEncoderStep,
+    MiniMaxH3Ref2VATextEncoderStep,
+    encode_vae_condition,
+)
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import MiniMaxH3DecodeStep
+from diffusers.modular_pipelines.minimax_h3.modular_pipeline import (
+    MiniMaxH3ModularPipeline,
+    align_num_frames,
+    audio_latent_num_frames,
+    resolve_canvas_size,
+    video_latent_num_frames,
+)
+from diffusers.modular_pipelines.modular_pipeline import (
+    BlockState,
+    ModularPipelineBlocks,
+    PipelineState,
+    SequentialPipelineBlocks,
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.utils import logging
+
+
+
+# ==============================================================================================================
+# Mask geometry — pixel-space masks onto the packed row grid
+#
+# A mask that selects *rows* has to be reduced on three grids, none of which a generic resize reproduces: the video
+# VAE's 16x spatial compression and then the transformer's 2x2 patch, because a row is one token; the VAE's chunked
+# causal frame grouping, which is not uniform; and, for audio, 40 latents per second rather than 24 frames per
+# second. Every reduction is a maximum — a row regenerates as much as the most-masked pixel it covers asks it to.
+# ==============================================================================================================
+
+
+logger = logging.get_logger(__name__)
+
+
+# The timesteps a fully preserved row is pinned at. MiniMax-H3 holds visual conditioning just short of clean and
+# audio conditioning exactly clean, and a preserved inpainting row is conditioning in every sense that matters, so it
+# is presented at the levels the checkpoint was trained to read.
+VISUAL_COND_TIMESTEP = 0.999
+AUDIO_COND_TIMESTEP = 1.0
+
+
+def resample_frame_indices(num_frames: int, fps: float, target_fps: float) -> np.ndarray:
+    r"""
+    The source frame every target-rate frame is taken from.
+
+    Whole frames are held and dropped, never blended, reproducing `ffmpeg`'s `fps` filter the way the reference
+    implementation resamples a reference clip. Returned as indices rather than frames so a clip and the mask drawn
+    over it go through the same resample.
+
+    Args:
+        num_frames (`int`): Frames the source carries.
+        fps (`float`): The rate it carries them at.
+        target_fps (`float`): The rate to resample onto, i.e. `components.fps`.
+
+    Returns:
+        `np.ndarray`: One source index per resampled frame.
+    """
+    if fps <= 0:
+        raise ValueError(f"A source clip must have a positive frame rate, got {fps}.")
+    if fps == target_fps:
+        return np.arange(num_frames)
+    scale = target_fps / fps
+    slots = np.floor(np.arange(num_frames) * scale + 0.5).astype(np.int64)
+    repeats = np.diff(slots, append=math.floor(num_frames * scale + 0.5))
+    return np.repeat(np.arange(num_frames), repeats)
+
+
+def normalize_source_frames(frames) -> np.ndarray:
+    r"""Any accepted clip layout as `uint8` `(num_frames, height, width, 3)`, matching the reference blocks."""
+    if isinstance(frames, list):
+        frames = np.stack([np.asarray(frame.convert("RGB")) for frame in frames])
+    if isinstance(frames, torch.Tensor):
+        frames = frames.movedim(-3, -1).cpu().numpy()
+    frames = np.asarray(frames)
+    if frames.dtype != np.uint8:
+        frames = (frames * 255.0).round().clip(0, 255).astype(np.uint8)
+    if frames.ndim != 4 or frames.shape[3] != 3:
+        raise ValueError(f"A source clip must be `(num_frames, height, width, 3)` RGB, got {tuple(frames.shape)}.")
+    return frames
+
+
+def normalize_source_mask(mask) -> np.ndarray:
+    r"""Any accepted mask layout as float32 `(num_frames, height, width)` over `[0, 1]`."""
+    if isinstance(mask, list):
+        mask = np.stack([np.asarray(frame.convert("L")) for frame in mask])
+    if isinstance(mask, torch.Tensor):
+        mask = mask.cpu().numpy()
+    mask = np.asarray(mask)
+    if mask.ndim == 4:
+        # `(num_frames, height, width, 1)` or a 3-channel mask painted as an image.
+        mask = mask.mean(axis=-1) if mask.shape[-1] in (3, 4) else mask[..., 0]
+    if mask.ndim == 2:
+        mask = mask[None]
+    if mask.ndim != 3:
+        raise ValueError(f"A mask must be `(num_frames, height, width)`, got {tuple(mask.shape)}.")
+    mask = mask.astype(np.float32)
+    if mask.max() > 1.0:
+        mask = mask / 255.0
+    return mask.clip(0.0, 1.0)
+
+
+def normalize_waveform(waveform: torch.Tensor, sample_rate: int, target_sample_rate: int) -> torch.Tensor:
+    r"""A `(channels, num_samples)` waveform as MiniMax-H3's stereo, float32, `target_sample_rate` soundtrack."""
+    import torchaudio
+
+    if waveform.ndim == 1:
+        waveform = waveform[None]
+    waveform = waveform.to(torch.float32)
+    if waveform.shape[0] == 1:
+        waveform = waveform.repeat(2, 1)
+    waveform = waveform[:2]
+    if sample_rate != target_sample_rate:
+        waveform = torchaudio.functional.resample(waveform, sample_rate, target_sample_rate)
+    return waveform
+
+
+class MiniMaxH3InpaintSourceGeometryStep(ModularPipelineBlocks):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Resolves the request's geometry from the source clip: the canvas its aspect ratio maps onto and the "
+            "`17 * n + 5` frame count its length rounds to. `ref2va` otherwise generates on MiniMax-H3's own 16:9 "
+            "canvas — references never bind the generated geometry — but an inpainting request has to land back on "
+            "top of its own footage, so the source is what settles it."
+        )
+
+    @property
+    def expected_configs(self):
+        from diffusers.modular_pipelines.modular_pipeline_utils import ConfigSpec
+
+        return [ConfigSpec("canvas_short_edge", 768), ConfigSpec("canvas_max_pixels", 768 * 1344)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="source_video",
+                required=True,
+                description=(
+                    "The clip being inpainted: a list of images, a `(num_frames, height, width, 3)` array or a "
+                    "`(num_frames, 3, height, width)` tensor, `uint8` or floating point over `[0, 1]`."
+                ),
+            ),
+            InputParam(
+                name="source_fps",
+                type_hint=float,
+                description="The rate `source_video` carries its frames at. Left out, MiniMax-H3's own 24 fps.",
+            ),
+            InputParam.template("height", description="Height of the generated video in pixels, a multiple of 32."),
+            InputParam.template("width", description="Width of the generated video in pixels, a multiple of 32."),
+            InputParam(
+                name="num_frames",
+                type_hint=int,
+                description="Frames to generate. Left out, as many as the source clip has at 24 fps.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
+            OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
+            OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5."),
+            OutputParam(
+                "source_frame_indices",
+                type_hint=np.ndarray,
+                description=(
+                    "The source frame every generated frame is taken from, after the 24 fps resample and the "
+                    "alignment padding. The clip and the mask drawn over it are both read through it."
+                ),
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        frames = normalize_source_frames(block_state.source_video)
+        fps = block_state.source_fps or components.fps
+
+        indices = resample_frame_indices(frames.shape[0], fps, components.fps)
+        num_frames = block_state.num_frames or int(indices.shape[0])
+
+        # The ceiling is a model limit, not a preference: past it the rotary clock leaves the trained range.
+        max_frames = int(components.max_duration * components.fps)
+        if num_frames > max_frames:
+            logger.warning(
+                f"MiniMax-H3 generates at most {components.max_duration:g} seconds; the source clip is "
+                f"{indices.shape[0] / components.fps:.2f}s at {components.fps} fps and is truncated to {max_frames} "
+                "frames. Inpaint a longer clip in segments."
+            )
+            num_frames = max_frames
+        num_frames = align_num_frames(num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk)
+        if num_frames < int(components.min_duration * components.fps):
+            raise ValueError(
+                f"MiniMax-H3 generates at least {components.min_duration:g} seconds, i.e. "
+                f"{int(components.min_duration * components.fps)} frames; the source clip resolves to {num_frames}."
+            )
+
+        # Alignment rounds *up*, so a clip that ends mid-chunk is held on its last frame rather than truncated down to
+        # the chunk below — the tail is a repeated frame the caller can trim, not lost footage.
+        if indices.shape[0] < num_frames:
+            indices = np.concatenate([indices, np.full(num_frames - indices.shape[0], indices[-1])])
+        block_state.source_frame_indices = indices[:num_frames]
+
+        if (block_state.height is None) != (block_state.width is None):
+            raise ValueError("`height` and `width` have to be passed together, or neither of them.")
+        if block_state.height is None:
+            block_state.height, block_state.width = resolve_canvas_size(
+                frames.shape[2],
+                frames.shape[1],
+                components.canvas_multiple,
+                components.config.canvas_short_edge,
+                components.config.canvas_max_pixels,
+            )
+        block_state.num_frames = num_frames
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3InpaintEncodeStep(ModularPipelineBlocks):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Encodes the source clip and its soundtrack, and reduces the masks onto the packed sequence's row grid. "
+            "The clip goes through the same recipe a keyframe anchor does — ImageNet normalization, a posterior "
+            "sampled under the fixed encode seed, float16 rounding, latent normalization — because that is what a "
+            "preserved row is presented as. The masks are reduced with a maximum on all three grids: the VAE's 16x "
+            "spatial compression, its `(1, 4, 4, 4, 4)` frame grouping, and the transformer's `2 x 2` patch."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("vae", AutoencoderKLMiniMaxH3),
+            ComponentSpec("audio_vae", AutoencoderKLMiniMaxH3Audio),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="source_video", required=True, description="The clip being inpainted. See above."),
+            InputParam(
+                name="source_frame_indices",
+                type_hint=np.ndarray,
+                required=True,
+                description="The source frame every generated frame is taken from.",
+            ),
+            InputParam(
+                name="mask",
+                required=True,
+                description=(
+                    "Which of the source clip's pixels the model repaints: `(num_frames, height, width)` over "
+                    "`[0, 1]`, `1` where the video regenerates and `0` where it is preserved, one frame per *source* "
+                    "frame. Read at its own resolution — it does not have to match the canvas, only the framing."
+                ),
+            ),
+            InputParam(
+                name="source_audio",
+                type_hint=torch.Tensor,
+                description=(
+                    "The source soundtrack, a `(channels, num_samples)` waveform. Left out, MiniMax-H3 writes the "
+                    "whole soundtrack itself."
+                ),
+            ),
+            InputParam(
+                name="source_audio_sample_rate",
+                type_hint=int,
+                description="The rate `source_audio` carries its samples at. Left out, the audio VAE's own.",
+            ),
+            InputParam(
+                name="audio_mask",
+                type_hint=torch.Tensor,
+                description=(
+                    "Which of the soundtrack the model rewrites, over the generated duration: `1` regenerates, `0` "
+                    "preserves. Any length, resampled onto the 40 Hz audio latent clock. Left out, a source "
+                    "soundtrack is preserved whole — which is what makes the model animate to the words already "
+                    "there rather than inventing new ones."
+                ),
+            ),
+            InputParam.template("height", required=True),
+            InputParam.template("width", required=True),
+            InputParam(name="num_frames", type_hint=int, required=True, description="Resolved number of frames."),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "source_latents",
+                type_hint=torch.Tensor,
+                description="The source clip encoded onto the target canvas, `(1, C, F, latent_h, latent_w)`.",
+            ),
+            OutputParam(
+                "source_audio_rows",
+                type_hint=torch.Tensor,
+                description="The source soundtrack in channel-major row layout, or None.",
+            ),
+            OutputParam(
+                "inpaint_row_mask",
+                type_hint=torch.Tensor,
+                description="One value per generated video row: `1` regenerates, `0` preserves.",
+            ),
+            OutputParam(
+                "inpaint_audio_row_mask",
+                type_hint=torch.Tensor,
+                description="One value per generated audio row, channel-major.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        height, width, num_frames = block_state.height, block_state.width, block_state.num_frames
+        ratio = components.vae_spatial_compression_ratio
+        num_latent_frames = video_latent_num_frames(
+            num_frames, components.vae_frames_per_chunk, components.vae_latents_per_chunk
+        )
+        num_audio_latents = audio_latent_num_frames(num_frames, components.fps)
+        indices = block_state.source_frame_indices
+
+        # 1. The clip, onto the 24 fps grid and the target canvas. LANCZOS is the reference implementation's rescale.
+        frames = normalize_source_frames(block_state.source_video)[indices]
+        if frames.shape[1:3] != (height, width):
+            frames = np.stack(
+                [
+                    np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS))
+                    for frame in frames
+                ]
+            )
+        pixels = torch.from_numpy(frames.copy()).to(device).permute(3, 0, 1, 2)[None]
+        block_state.source_latents = encode_vae_condition(
+            components.vae,
+            pixels,
+            components.pixel_mean,
+            components.pixel_std,
+            components.keyframe_encode_seed,
+        )
+
+        # 2. The mask, on the same 24 fps grid but at its own resolution — the reduction onto the latent grid is a
+        # maximum, so resizing it to the canvas first would only cost detail.
+        mask = torch.from_numpy(normalize_source_mask(block_state.mask)[indices].copy())
+        block_state.inpaint_row_mask = quantize_mask(
+            pixel_mask_to_row_mask(mask, num_latent_frames, height // ratio, width // ratio, components.patch_size)
+        ).to(device)
+
+        # 3. The soundtrack. A request without one generates its audio outright, so every audio row is masked in.
+        num_audio_rows = num_audio_latents * components.audio_channels
+        if block_state.source_audio is None:
+            block_state.source_audio_rows = None
+            block_state.inpaint_audio_row_mask = torch.ones(num_audio_rows, device=device)
+        else:
+            sample_rate = block_state.source_audio_sample_rate or components.audio_sampling_rate
+            waveform = normalize_waveform(block_state.source_audio, sample_rate, components.audio_sampling_rate)
+            # The audio VAE right-pads to a whole hop, so exactly `num_audio_latents` hops encode to exactly the rows
+            # the layout reserves. Silence pads a soundtrack shorter than the video.
+            num_samples = num_audio_latents * components.audio_vae.hop_length
+            if waveform.shape[-1] < num_samples:
+                waveform = F.pad(waveform, (0, num_samples - waveform.shape[-1]))
+            waveform = waveform[:, :num_samples]
+
+            # Under `native` attention explicitly: the audio VAE's encoder carries a causal-attention `pre_block` kept
+            # in float32, and cuDNN has no float32 SDPA kernel, so the backend a pipeline selects for its transformer
+            # aborts here. Only *encoding* audio reaches this path.
+            with attention_backend("native"):
+                posterior = components.audio_vae.encode(waveform.to(device)[:, None], return_dict=False)[0]
+            latents = posterior.mode().float().transpose(1, 2)
+            mean = torch.tensor(components.audio_vae.config.latents_mean, device=latents.device).view(1, 1, -1)
+            std = torch.tensor(components.audio_vae.config.latents_std, device=latents.device).view(1, 1, -1)
+            rows = ((latents - mean) / std).reshape(-1, components.audio_latent_channels)
+            if rows.shape[0] != num_audio_rows:
+                raise ValueError(
+                    f"The source soundtrack encodes to {rows.shape[0]} audio rows but the layout holds "
+                    f"{num_audio_rows}."
+                )
+            block_state.source_audio_rows = rows
+
+            if block_state.audio_mask is None:
+                block_state.inpaint_audio_row_mask = torch.zeros(num_audio_rows, device=device)
+            else:
+                block_state.inpaint_audio_row_mask = quantize_mask(
+                    audio_mask_to_row_mask(
+                        block_state.audio_mask, num_audio_latents, components.audio_channels
+                    )
+                ).to(device)
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3InpaintPrepareStep(ModularPipelineBlocks):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Packs the source latents into rows, checks them against the layout, keeps the noise the generated rows "
+            "were drawn from, and imposes the source once before the first forward. From here the loop reads a "
+            "sequence whose preserved rows already hold the source at the conditioning level."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="source_latents", type_hint=torch.Tensor, required=True, description="The encoded source clip."
+            ),
+            InputParam(
+                name="source_audio_rows",
+                type_hint=torch.Tensor,
+                description="The encoded source soundtrack in row layout, or None.",
+            ),
+            InputParam(
+                name="inpaint_row_mask", type_hint=torch.Tensor, required=True, description="Per video row."
+            ),
+            InputParam(
+                name="inpaint_audio_row_mask", type_hint=torch.Tensor, required=True, description="Per audio row."
+            ),
+            InputParam(
+                name="latents",
+                type_hint=torch.Tensor,
+                required=True,
+                description="The video rows of the packed sequence, conditioning rows first.",
+            ),
+            InputParam(
+                name="audio_latents",
+                type_hint=torch.Tensor,
+                required=True,
+                description="The audio rows of the packed sequence, reference rows first.",
+            ),
+            InputParam(
+                name="num_condition_video_rows",
+                type_hint=int,
+                default=0,
+                description="How many leading video rows are conditioning rows.",
+            ),
+            InputParam(
+                name="num_condition_audio_rows",
+                type_hint=int,
+                default=0,
+                description="How many leading audio rows are reference rows.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("latents", type_hint=torch.Tensor, description="The video rows, source imposed."),
+            OutputParam("audio_latents", type_hint=torch.Tensor, description="The audio rows, source imposed."),
+            OutputParam("source_rows", type_hint=torch.Tensor, description="The source clip in row layout."),
+            OutputParam(
+                "inpaint_noise_rows",
+                type_hint=torch.Tensor,
+                description="The noise the generated video rows were drawn from.",
+            ),
+            OutputParam(
+                "inpaint_audio_noise_rows",
+                type_hint=torch.Tensor,
+                description="The noise the generated audio rows were drawn from.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        num_condition_video_rows = block_state.num_condition_video_rows
+        num_condition_audio_rows = block_state.num_condition_audio_rows
+
+        source_rows = patchify_video_latents(
+            block_state.source_latents.to(device, torch.float32), components.patch_size
+        )
+        generated_video_rows = block_state.latents.shape[0] - num_condition_video_rows
+        if source_rows.shape[0] != generated_video_rows:
+            raise ValueError(
+                f"The source clip packs into {source_rows.shape[0]} video rows but the layout reserved "
+                f"{generated_video_rows} generated ones. The canvas the layout was built from and the one the source "
+                "was encoded at do not agree."
+            )
+        if block_state.inpaint_row_mask.shape[0] != generated_video_rows:
+            raise ValueError(
+                f"The mask reduces to {block_state.inpaint_row_mask.shape[0]} rows but the layout reserved "
+                f"{generated_video_rows} generated video rows."
+            )
+        block_state.source_rows = source_rows
+
+        block_state.inpaint_noise_rows = block_state.latents[num_condition_video_rows:].clone()
+        block_state.inpaint_audio_noise_rows = block_state.audio_latents[num_condition_audio_rows:].clone()
+
+        # The first forward reads the same presentation every later one does.
+        block_state.latents[num_condition_video_rows:] = impose_source(
+            block_state.latents[num_condition_video_rows:],
+            source_rows,
+            block_state.inpaint_noise_rows,
+            block_state.inpaint_row_mask,
+            VISUAL_COND_TIMESTEP,
+            components.scheduler,
+        )
+        if block_state.source_audio_rows is not None:
+            block_state.audio_latents[num_condition_audio_rows:] = impose_source(
+                block_state.audio_latents[num_condition_audio_rows:],
+                block_state.source_audio_rows.to(device, torch.float32),
+                block_state.inpaint_audio_noise_rows,
+                block_state.inpaint_audio_row_mask,
+                AUDIO_COND_TIMESTEP,
+                components.audio_scheduler,
+            )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+def impose_source(
+    rows: torch.Tensor,
+    source_rows: torch.Tensor,
+    noise_rows: torch.Tensor,
+    row_mask: torch.Tensor,
+    condition_timestep: float,
+    scheduler,
+) -> torch.Tensor:
+    r"""
+    Blend the source into the rows the mask preserves.
+
+    The preserved content enters at the *conditioning* timestep rather than the step's own — clean for audio, `0.999`
+    for video — and the mask does the mixing, so a row with mask `m` ends up carrying content at roughly
+    `1 - m * sigma`, which is exactly the timestep the layout labels it with.
+
+    Args:
+        rows (`torch.Tensor` of shape `(num_rows, dim)`): The generated rows as the loop last left them.
+        source_rows (`torch.Tensor` of shape `(num_rows, dim)`): The source in the same layout.
+        noise_rows (`torch.Tensor` of shape `(num_rows, dim)`): The noise those rows were drawn from.
+        row_mask (`torch.Tensor` of shape `(num_rows,)`): `1` regenerates, `0` preserves.
+        condition_timestep (`float`): The level the source is presented at.
+        scheduler (`MiniMaxH3Scheduler`): The stream's scheduler, for its forward process.
+
+    Returns:
+        `torch.Tensor`: The rows with the source imposed.
+    """
+    mask = row_mask.to(rows.dtype).unsqueeze(-1)
+    source = source_rows.to(rows)
+    if condition_timestep < 1.0:
+        source = scheduler.scale_noise(source, condition_timestep, noise_rows.to(rows))
+    return mask * rows + (1.0 - mask) * source
+
+
+def _place_rows(row_mask: torch.Tensor, timestep: float, pin: float) -> torch.Tensor:
+    r"""
+    Place every generated row of one stream on its own timestep: `1 - m * sigma`, pinned at `pin`.
+
+    The two ends are *selected* rather than computed. Every distinct value here becomes a row of the transformer's
+    modulation table, and float32 arithmetic lands `1 - 1 * (1 - t)` a ulp away from `t` — which would split each
+    level in two and leave the fully generated rows claiming a timestep the scheduler is not actually stepping them
+    on. Computing in float64 and selecting the endpoints keeps a hard mask exactly as cheap as no mask at all.
+
+    Args:
+        row_mask (`torch.Tensor` of shape `(num_rows,)`): `1` regenerates, `0` preserves.
+        timestep (`float`): The stream's own timestep this step.
+        pin (`float`): Where a fully preserved row sits, i.e. the conditioning timestep or the schedule if it is
+            already cleaner.
+
+    Returns:
+        `torch.Tensor` of shape `(num_rows,)`: the float32 timestep of every row.
+    """
+    mask = row_mask.detach().cpu().to(torch.float64)
+    rows = (1.0 - mask * (1.0 - timestep)).clamp(max=pin)
+    rows = torch.where(mask >= 1.0, torch.tensor(timestep, dtype=torch.float64), rows)
+    rows = torch.where(mask <= 0.0, torch.tensor(pin, dtype=torch.float64), rows)
+    return rows.to(torch.float32)
+
+
+class MiniMaxH3InpaintSetTimestepsStep(MiniMaxH3SetTimestepsStep):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "The `ref2va` timestep plan, with every generated row relabelled by the mask. A row masked `m` sits at "
+            "`1 - m * sigma` of its own stream's schedule, clamped at the conditioning timestep: fully generated rows "
+            "keep the schedule, fully preserved ones land on `0.999` (video) or `1.0` (audio) — where the checkpoint "
+            "reads clean conditioning — and feathered ones fall honestly in between. MiniMax-H3 already carries a "
+            "per-row timestep vector for its anchors, so nothing about the transformer changes."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(
+                name="inpaint_row_mask", type_hint=torch.Tensor, required=True, description="Per video row."
+            ),
+            InputParam(
+                name="inpaint_audio_row_mask", type_hint=torch.Tensor, required=True, description="Per audio row."
+            ),
+        ]
+
+    @staticmethod
+    def build_masked_row_timesteps(
+        video_indices: torch.Tensor,
+        audio_indices: torch.Tensor,
+        num_condition_video_rows: int,
+        num_condition_audio_rows: int,
+        num_text_tokens: int,
+        video_timestep: float,
+        audio_timestep: float,
+        condition_video_timestep: float,
+        condition_audio_timestep: float,
+        video_row_mask: torch.Tensor,
+        audio_row_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        [`MiniMaxH3SetTimestepsStep.build_row_timesteps`] with the generated rows placed by the mask.
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor]`: the distinct timesteps, sorted, and the index of every row into them.
+        """
+        sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens)
+        row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)
+
+        # `1 - t` is the stream's sigma; a row only carries `m` of it, and a fully preserved row is pinned wherever
+        # the checkpoint reads conditioning — which the schedule itself overtakes on the last steps.
+        video_rows = _place_rows(video_row_mask, video_timestep, max(video_timestep, condition_video_timestep))
+        audio_rows = _place_rows(audio_row_mask, audio_timestep, max(audio_timestep, condition_audio_timestep))
+
+        row_timesteps[video_indices[num_condition_video_rows:]] = video_rows
+        row_timesteps[video_indices[:num_condition_video_rows]] = condition_video_timestep
+        row_timesteps[audio_indices[num_condition_audio_rows:]] = audio_rows
+        row_timesteps[audio_indices[:num_condition_audio_rows]] = condition_audio_timestep
+        return torch.unique(row_timesteps, sorted=True, return_inverse=True)
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        components.audio_scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        block_state.timesteps = components.scheduler.timesteps
+        block_state.audio_timesteps = components.audio_scheduler.timesteps
+
+        block_state.row_timestep_plan = [
+            tuple(
+                tensor.to(device)
+                for tensor in self.build_masked_row_timesteps(
+                    block_state.video_indices,
+                    block_state.audio_indices,
+                    block_state.num_condition_video_rows,
+                    block_state.num_condition_audio_rows,
+                    block_state.text_indices.numel(),
+                    float(timestep),
+                    float(audio_timestep),
+                    max(float(timestep), components.keyframe_noise_aug),
+                    AUDIO_COND_TIMESTEP,
+                    block_state.inpaint_row_mask,
+                    block_state.inpaint_audio_row_mask,
+                )
+            )
+            for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+        ]
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MiniMaxH3InpaintLoopSchedulerStep(MiniMaxH3LoopSchedulerStep):
+    model_name = "minimax-h3"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Steps both streams, then re-imposes the source on every row the mask preserves. The scheduler is still "
+            "stepped over those rows so its step index keeps up with the loop; the write-back is what the next "
+            "forward reads. The last step imposes the source clean, so what the mask preserved decodes back to the "
+            "footage it came from."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(name="source_rows", type_hint=torch.Tensor, required=True, description="The source clip."),
+            InputParam(
+                name="source_audio_rows", type_hint=torch.Tensor, description="The source soundtrack, or None."
+            ),
+            InputParam(
+                name="inpaint_noise_rows", type_hint=torch.Tensor, required=True, description="Video row noise."
+            ),
+            InputParam(
+                name="inpaint_audio_noise_rows", type_hint=torch.Tensor, required=True, description="Audio row noise."
+            ),
+            InputParam(
+                name="inpaint_row_mask", type_hint=torch.Tensor, required=True, description="Per video row."
+            ),
+            InputParam(
+                name="inpaint_audio_row_mask", type_hint=torch.Tensor, required=True, description="Per audio row."
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        components, block_state = super().__call__(components, block_state, i, t)
+
+        # On the last step the schedule has reached the clean end, and so does the source: the preserved rows are
+        # handed to the decoder as the footage itself rather than as the anchor's 0.999.
+        last = i + 1 >= block_state.timesteps.numel()
+        video_level = 1.0 if last else VISUAL_COND_TIMESTEP
+
+        num_condition_video_rows = block_state.num_condition_video_rows
+        block_state.latents[num_condition_video_rows:] = impose_source(
+            block_state.latents[num_condition_video_rows:],
+            block_state.source_rows,
+            block_state.inpaint_noise_rows,
+            block_state.inpaint_row_mask,
+            video_level,
+            components.scheduler,
+        )
+        if block_state.source_audio_rows is not None:
+            num_condition_audio_rows = block_state.num_condition_audio_rows
+            block_state.audio_latents[num_condition_audio_rows:] = impose_source(
+                block_state.audio_latents[num_condition_audio_rows:],
+                block_state.source_audio_rows,
+                block_state.inpaint_audio_noise_rows,
+                block_state.inpaint_audio_row_mask,
+                AUDIO_COND_TIMESTEP,
+                components.audio_scheduler,
+            )
+        return components, block_state
+
+
+class MiniMaxH3InpaintDenoiseStep(MiniMaxH3DenoiseLoopWrapper):
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3Ref2VALoopDenoiser, MiniMaxH3InpaintLoopSchedulerStep]
+    block_names = ["denoiser", "update"]
+
+    @property
+    def description(self) -> str:
+        return "Runs the `ref2va` denoising loop with the source imposed on every preserved row, one forward per step."
+
+
+class MiniMaxH3Ref2VAInpaintCoreDenoiseStep(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3Ref2VAPrepareLayoutStep,
+        MiniMaxH3PrepareConditionLatentsStep,
+        MiniMaxH3PrepareLatentsStep,
+        MiniMaxH3Ref2VAPrepareLatentsStep,
+        MiniMaxH3InpaintPrepareStep,
+        MiniMaxH3InpaintSetTimestepsStep,
+        MiniMaxH3InpaintDenoiseStep,
+        MiniMaxH3AfterDenoiseStep,
+    ]
+    block_names = [
+        "prepare_layout",
+        "prepare_condition_latents",
+        "prepare_latents",
+        "prepare_latents_ref2va",
+        "prepare_inpaint",
+        "set_timesteps",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Core denoising workflow for masked `ref2va` inpainting: `MiniMaxH3Ref2VACoreDenoiseStep` with the source "
+            "clip imposed on every row the mask preserves, and those rows relabelled onto the conditioning timestep."
+        )
+
+
+class MiniMaxH3Ref2VAInpaintBlocks(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    block_classes = [
+        MiniMaxH3InpaintSourceGeometryStep,
+        MiniMaxH3Ref2VASetupStep,
+        MiniMaxH3Ref2VATextEncoderStep,
+        MiniMaxH3Ref2VAReferenceEncoderStep,
+        MiniMaxH3InpaintEncodeStep,
+        MiniMaxH3Ref2VAInpaintCoreDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = [
+        "source_geometry",
+        "setup",
+        "text_encoder",
+        "reference_encoder",
+        "inpaint_encoder",
+        "denoise",
+        "decode",
+    ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Masked video and audio inpainting with MiniMax-H3, on the `ref2va` reference workflow: the source clip "
+            "settles the canvas and the frame count, the references condition the repaint the way they condition any "
+            "`ref2va` request, and the mask decides which rows the model writes. Everything the mask preserves rides "
+            "through the loop as conditioning, at the levels the checkpoint reads it at."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("videos", description="The inpainted video."),
+            OutputParam(
+                "audio",
+                type_hint=torch.Tensor,
+                description="The soundtrack of the packed sequence, of shape `(1, 2, num_samples)`.",
+            ),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the soundtrack in Hz."),
+        ]
+
+
+class MiniMaxH3Ref2VAInpaintGeneratorBlocks(SequentialPipelineBlocks):
+    model_name = "minimax-h3"
+    # `MiniMaxH3Ref2VAInpaintBlocks` without its text-encoder step, so the 62 GiB Qwen3-VL can live somewhere else and
+    # `prompt_embeds` / `text_token_tags` are the whole wire format — the split the reference and audio-to-video
+    # Spaces already deploy behind. The geometry step stays: it is also what maps source frames onto the 24 fps grid,
+    # and passing the conditioner's resolved `height` / `width` / `num_frames` through it leaves them untouched.
+    block_classes = [
+        MiniMaxH3InpaintSourceGeometryStep,
+        MiniMaxH3Ref2VASetupStep,
+        MiniMaxH3Ref2VAReferenceEncoderStep,
+        MiniMaxH3InpaintEncodeStep,
+        MiniMaxH3Ref2VAInpaintCoreDenoiseStep,
+        MiniMaxH3DecodeStep,
+    ]
+    block_names = ["source_geometry", "setup", "reference_encoder", "inpaint_encoder", "denoise", "decode"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "The denoising half of a split MiniMax-H3 deployment, masked: `MiniMaxH3Ref2VAInpaintBlocks` without its "
+            "text-encoder step. The caller has to pass the `height`, `width` and `num_frames` the conditioner "
+            "resolved, because the presentation the embeddings were built from encodes that plan."
+        )
+
+    @property
+    def outputs(self):
+        return MiniMaxH3Ref2VAInpaintBlocks.outputs.fget(self)
+
+# The VAE's chunked causal grouping, as pixel frames per latent frame within one 17-frame chunk. `clip_length = 17`
+# pre-pads to 20 at `vae_ratio_t = 4` and drops the 3 leading tokens, so the chunk's first kept latent frame covers
+# one real pixel frame and the remaining four cover four each.
+MINIMAX_H3_FRAMES_PER_CHUNK = 17
+MINIMAX_H3_FRAMES_PER_LATENT = (1, 4, 4, 4, 4)
+
+
+def latent_frame_groups(
+    num_frames: int,
+    num_latent_frames: int,
+    frames_per_chunk: int = MINIMAX_H3_FRAMES_PER_CHUNK,
+    frames_per_latent: tuple[int, ...] = MINIMAX_H3_FRAMES_PER_LATENT,
+) -> list[tuple[int, int]]:
+    r"""
+    The pixel frames each latent frame is encoded from.
+
+    Args:
+        num_frames (`int`): Pixel frames of the clip, of the form `17 * n + 5`.
+        num_latent_frames (`int`): Latent frames the VAE produces for them, `5 * n + 2`.
+        frames_per_chunk (`int`, defaults to 17): Pixel frames per VAE chunk, its `clip_length`.
+        frames_per_latent (`tuple[int, ...]`, defaults to `(1, 4, 4, 4, 4)`):
+            Pixel frames each latent frame of a chunk covers, cycling from the first frame.
+
+    Returns:
+        `list[tuple[int, int]]`: one `(start, end)` half-open pixel-frame range per latent frame, covering the clip.
+    """
+    starts = [0]
+    for span in frames_per_latent:
+        starts.append(starts[-1] + span)
+    cycle = len(frames_per_latent)
+
+    groups = []
+    for index in range(num_latent_frames):
+        start = (index // cycle) * frames_per_chunk + starts[index % cycle]
+        end = start + frames_per_latent[index % cycle]
+        groups.append((min(start, num_frames - 1), min(end, num_frames)))
+    # A frame count that ends mid-chunk leaves a tail no latent frame's own span reaches; it is encoded into the last
+    # latent frame, so that is where its coverage belongs.
+    groups[-1] = (groups[-1][0], num_frames)
+    return groups
+
+
+def pixel_mask_to_row_mask(
+    mask: torch.Tensor,
+    num_latent_frames: int,
+    latent_height: int,
+    latent_width: int,
+    patch_size: tuple[int, int, int] = (1, 2, 2),
+    frames_per_chunk: int = MINIMAX_H3_FRAMES_PER_CHUNK,
+    frames_per_latent: tuple[int, ...] = MINIMAX_H3_FRAMES_PER_LATENT,
+) -> torch.Tensor:
+    r"""
+    Reduce a pixel-space video mask to one value per video row of the packed sequence.
+
+    Every reduction is a maximum: a row regenerates as much as the most-masked pixel it covers asks it to. That is the
+    safe direction — a token the model must repaint is never accidentally pinned to the source — and it is what makes
+    a feathered mask behave, since the softest values survive into the row rather than being averaged away.
+
+    Args:
+        mask (`torch.Tensor` of shape `(num_frames, height, width)`):
+            The mask over the source clip, `1` where the video regenerates and `0` where it is preserved. Values in
+            between are honoured: they place the row part-way down its own schedule.
+        num_latent_frames (`int`), latent_height (`int`), latent_width (`int`):
+            The generated video's latent shape, i.e. what the layout step resolved.
+        patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): The transformer's `(t, h, w)` patch.
+        frames_per_chunk (`int`), frames_per_latent (`tuple[int, ...]`): See [`latent_frame_groups`].
+
+    Returns:
+        `torch.Tensor` of shape `(num_latent_frames * (latent_height // patch_h) * (latent_width // patch_w),)`:
+        one value per video row, in the frame-major then row-major order [`patchify_video_latents`] produces.
+    """
+    if mask.ndim != 3:
+        raise ValueError(f"A video mask must be `(num_frames, height, width)`, got {tuple(mask.shape)}.")
+    _, patch_h, patch_w = patch_size
+    if latent_height % patch_h or latent_width % patch_w:
+        raise ValueError(
+            f"A {latent_height}x{latent_width} latent canvas is not divisible by the patch {(patch_h, patch_w)}."
+        )
+
+    num_frames = mask.shape[0]
+    mask = mask.to(torch.float32).clamp(0.0, 1.0)
+
+    # 1. Spatially, onto the latent grid. `adaptive_max_pool2d` divides each axis into `latent_*` near-equal bands,
+    # which is the VAE's own 16x split whenever the canvas is a multiple of 16 — and it stays sane when it is not.
+    reduced = F.adaptive_max_pool2d(mask[:, None], (latent_height, latent_width))[:, 0]
+
+    # 2. Temporally, onto the VAE's chunked grouping.
+    groups = latent_frame_groups(num_frames, num_latent_frames, frames_per_chunk, frames_per_latent)
+    reduced = torch.stack([reduced[start:end].amax(dim=0) for start, end in groups])
+
+    # 3. Onto the transformer's patch. A row is one token: it carries one timestep and is written back as a whole, so
+    # sub-patch detail cannot survive and the strongest value in the patch is what the row acts on.
+    rows = reduced.reshape(
+        num_latent_frames, latent_height // patch_h, patch_h, latent_width // patch_w, patch_w
+    ).amax(dim=(2, 4))
+    return rows.reshape(-1)
+
+
+def audio_mask_to_row_mask(
+    mask: torch.Tensor,
+    num_audio_latents: int,
+    audio_channels: int = 2,
+) -> torch.Tensor:
+    r"""
+    Reduce a mask over the soundtrack's timeline to one value per audio row of the packed sequence.
+
+    Args:
+        mask (`torch.Tensor` of shape `(num_audio_latents,)` or `(n,)`):
+            The mask over the generated soundtrack, `1` where the audio regenerates and `0` where it is preserved. A
+            mask of another length is resampled onto the audio latent clock with a maximum, so a mask drawn at video
+            frame rate — or as a timeline image — lands on the right latents rather than half a beat off.
+        num_audio_latents (`int`): Audio latents per channel, i.e. what the layout step resolved.
+        audio_channels (`int`, defaults to 2): Channels the soundtrack is packed channel-major over.
+
+    Returns:
+        `torch.Tensor` of shape `(num_audio_latents * audio_channels,)`: one value per audio row, channel-major.
+    """
+    mask = mask.reshape(-1).to(torch.float32).clamp(0.0, 1.0)
+    if mask.shape[0] != num_audio_latents:
+        mask = F.adaptive_max_pool1d(mask[None, None], num_audio_latents)[0, 0]
+    # Channel-major: the layout lays both stereo channels out as two blocks of `num_audio_latents` rows, so one
+    # timeline repeats rather than interleaves.
+    return mask.repeat(audio_channels)
+
+
+def quantize_mask(mask: torch.Tensor, levels: int = 256) -> torch.Tensor:
+    r"""
+    Snap a mask to a grid of `levels` steps, rounding *up* so a partly-masked row never rounds to fully preserved.
+
+    Each distinct mask value becomes a distinct row timestep, and every distinct row timestep is an extra row in the
+    transformer's modulation table. A feathered mask left at float32 can carry thousands of them; on this grid it
+    carries at most `levels + 1`, and the difference is invisible at the noise levels involved.
+
+    Args:
+        mask (`torch.Tensor`): The mask to snap.
+        levels (`int`, defaults to 256): Steps to snap to.
+
+    Returns:
+        `torch.Tensor`: The snapped mask.
+    """
+    return torch.ceil(mask * levels) / levels

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "MiniMaxH3Ref2VAInpaintBlocks",
+  "_diffusers_version": "0.40.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "minimax_h3_inpaint_blocks.MiniMaxH3Ref2VAInpaintBlocks"
+  }
+}