minimax_h3_video ¶

def decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
    """Stream decoded ``[0, 1]`` FP32 pixels into a caller-owned CPU buffer."""
    expected_shape = self.decoded_pixel_shape(z.shape)
    if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
        raise ValueError(
            "`output` must be a CPU float32 tensor with shape "
            f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
    try:
        if self.use_slicing and z.shape[0] > 1:
            for batch_index, z_slice in enumerate(z.split(1)):
                self._decode_to_pixels(z_slice, output[batch_index:batch_index + 1])
        else:
            self._decode_to_pixels(z, output)
    finally:
        # Drain async chunk copies before the caller (or an exception
        # handler) can read or release the pinned buffer.
        if self._streams_chunk_copies(z, output):
            torch.cuda.current_stream(z.device).synchronize()
    return output

def decoded_pixel_shape(self, latent_shape: torch.Size | tuple[int, ...]) -> tuple[int, int, int, int, int]:
    """Return the exact CPU pixel-buffer shape for a latent tensor shape."""
    if len(latent_shape) != 5:
        raise ValueError(f"MiniMax-H3 latents must be five-dimensional, got shape {tuple(latent_shape)}.")
    batch_size, channels, latent_num_frames, latent_height, latent_width = map(int, latent_shape)
    if channels != self.latent_channels:
        raise ValueError(f"MiniMax-H3 latents must have {self.latent_channels} channels, got {channels}.")
    _, _, decoded_num_frames = self._temporal_decode_plan(latent_num_frames)
    return (
        batch_size,
        int(self.config.out_channels),
        decoded_num_frames,
        latent_height * self.spatial_compression_ratio,
        latent_width * self.spatial_compression_ratio,
    )

def encode_keyframe(
    self,
    x: torch.Tensor,
    return_dict: bool = True,
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
    """Encode one-frame conditioning inputs without video chunk padding."""
    if x.ndim != 5 or x.shape[2] != 1:
        raise ValueError(f"`x` must contain exactly one video frame, got shape {tuple(x.shape)}.")
    if self.use_slicing and x.shape[0] > 1:
        moments = torch.cat([self._encode_clip(x_slice) for x_slice in x.split(1)])
    else:
        moments = self._encode_clip(x)
    posterior = DiagonalGaussianDistribution(moments)
    if not return_dict:
        return (posterior, )
    return AutoencoderKLOutput(latent_dist=posterior)

def encode_pixels(
    self,
    pixels: torch.Tensor,
    return_dict: bool = True,
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
    """Encode CPU-resident pixels one VAE clip at a time.

    ``pixels`` stays on CPU as ``uint8`` in ``[0, 255]`` or floating point
    in ``[0, 1]``; each clip is moved to the VAE device, normalized, and
    encoded so only one clip of pixels is resident on the accelerator.
    """
    if pixels.ndim != 5 or pixels.shape[1] != self.config.in_channels or pixels.shape[2] <= 0:
        raise ValueError(
            f"`pixels` must have shape [B, {self.config.in_channels}, T, H, W] with T > 0, "
            f"got {tuple(pixels.shape)}.")
    if pixels.device.type != "cpu":
        raise ValueError(f"`pixels` must remain on CPU, got device={pixels.device}.")
    if pixels.dtype != torch.uint8 and not pixels.is_floating_point():
        raise TypeError(f"`pixels` must use uint8 or a floating-point dtype, got {pixels.dtype}.")
    if self.use_slicing and pixels.shape[0] > 1:
        moments = torch.cat([self._encode_pixels(pixel_slice) for pixel_slice in pixels.split(1)])
    else:
        moments = self._encode_pixels(pixels)
    posterior = DiagonalGaussianDistribution(moments)
    if not return_dict:
        return (posterior, )
    return AutoencoderKLOutput(latent_dist=posterior)

def prepare_for_compile(self) -> None:
    """Compile the fixed-shape tile helpers for the opt-in VAE compile path.

    ``ComposedPipelineBase._maybe_compile_pipeline_module`` calls this hook
    only when ``enable_torch_compile_vae`` is set, right before the decoder
    is compiled through ``_compile_conditions``. The spatial tile grid and
    the per-tile decoder-input projection have fixed shapes, so
    ``mode="reduce-overhead"`` records one CUDA graph per geometry and
    replays it for every tile and temporal chunk. Keeping this behind the
    opt-in means default (eager) users pay neither the inductor/triton
    toolchain requirement and first-decode compile latency nor the
    permanent cudagraph memory pools, and multi-resolution callers never
    churn ``dynamic=False`` recompiles they did not ask for.
    """
    if self._tile_helpers_compiled:
        return
    # The fixed spatial tile grid reuses one compiled blend-and-concatenate graph.
    self._stitch_tiles = torch.compile(self._stitch_tiles, backend="inductor", mode="reduce-overhead", dynamic=False)
    # Each fixed-shape latent tile reuses one compiled decoder-input projection.
    self._project_decoder_tile = torch.compile(
        self._project_decoder_tile,
        backend="inductor",
        mode="reduce-overhead",
        dynamic=False,
    )
    self._tile_helpers_compiled = True

def forward(
    self,
    hidden_states: torch.Tensor,
    rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Apply dense self-attention to one spatial VAE token sequence."""
    query = self.to_q(hidden_states).unflatten(2, (self.heads, -1))
    key = self.to_k(hidden_states).unflatten(2, (self.heads, -1))
    value = self.to_v(hidden_states).unflatten(2, (self.heads, -1))

    query = self.norm_q(query.float()).to(query.dtype)
    key = self.norm_k(key.float()).to(key.dtype)
    if rotary_emb is not None:
        cos, sin = rotary_emb
        cos = cos.to(query.dtype)
        sin = sin.to(query.dtype)
        rotary_dim = cos.shape[-1]
        query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
        key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
        query_first, query_second = query_rotary.chunk(2, dim=-1)
        key_first, key_second = key_rotary.chunk(2, dim=-1)
        query_rotated = torch.cat([-query_second, query_first], dim=-1)
        key_rotated = torch.cat([-key_second, key_first], dim=-1)
        query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1)
        key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1)

    if self.attn_impl is not None and query.device.type != "cpu":
        # VAE decoding has no diffusion-step metadata, so call the selected
        # backend implementation directly with dense BSHD tensors.
        hidden_states = self.attn_impl.forward(query, key, value, None)
        hidden_states = hidden_states.flatten(2, 3)
    else:
        # Keep CPU construction and execution available without requiring
        # an accelerator attention backend.
        query, key, value = (tensor.permute(0, 2, 1, 3) for tensor in (query, key, value))
        hidden_states = F.scaled_dot_product_attention(query, key, value)
        hidden_states = hidden_states.permute(0, 2, 1, 3).flatten(2, 3)
    return self.to_out[0](hidden_states)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Decode one latent spatial input through the H3 video transformer."""
    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(
        batch_size,
        num_frames * height * width,
        num_channels,
    )
    hidden_states = self.proj_in(hidden_states)
    num_patches = hidden_states.shape[1]
    register_tokens = self.register_tokens.expand(batch_size, -1, -1)
    cls_token = torch.zeros_like(hidden_states[:, :1, :])
    hidden_states = torch.cat([hidden_states, register_tokens, cls_token], dim=1)

    grids = [
        2.0 * (torch.arange(0.5, size, dtype=torch.float32, device=hidden_states.device) / size) - 1.0
        for size in (num_frames, height, width)
    ]
    position_ids = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=-1).flatten(0, 2)
    position_ids = position_ids.unsqueeze(0).expand(batch_size, -1, -1)
    suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3))
    rotary_emb = self.rope(torch.cat([position_ids, suffix_ids], dim=1))

    for block in self.transformer_blocks:
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            hidden_states = checkpoint(block, hidden_states, rotary_emb, use_reentrant=False)
        else:
            hidden_states = block(hidden_states, rotary_emb)

    hidden_states = self.proj_out(self.norm_out(hidden_states))[:, :num_patches, :]
    patch_size, patch_size_t = self.patch_size, self.patch_size_t
    hidden_states = hidden_states.view(
        batch_size,
        num_frames,
        height,
        width,
        self.out_channels,
        patch_size_t,
        patch_size,
        patch_size,
    )
    hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
    return hidden_states.reshape(
        batch_size,
        self.out_channels,
        num_frames * patch_size_t,
        height * patch_size,
        width * patch_size,
    )