Compare commits

...
Author SHA1 Message Date
Will LinandClaude Fable 5 942f7db3db [test] H3 VAE streaming: legacy-decode oracle, batch slicing, pinned-buffer coverage
Pin the chunk-iterator refactor to the pre-streaming _decode implementation
bit-for-bit across the seam and pad-trim geometries (one padded chunk, a pad
hitting the intra-clip tail, two blended chunks, three chunks plus trim), and
cover the use_slicing batch paths for encode_pixels/decode_to_pixels and the
CUDA pinned-buffer async copy path (skipped without a GPU).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-21 02:38:44 +00:00
Will LinandClaude Fable 5 aadb23f409 [perf] H3 VAE streamed decode: direct per-plane copies, async with pinned buffers
The temporal slice of the CPU output buffer is strided across channels, so
each finalized-chunk copy_ staged through a pageable CPU temporary plus a
CPU-side scatter, which is where the streamed path's decode-time regression
came from and why the pinned buffer bought nothing. Copy per (batch, channel)
plane instead - contiguous on both sides, memcpy-eligible - and make the
copies non_blocking when the destination is pinned, draining the stream once
in decode_to_pixels before the buffer can be read or released.

Also: raise on a non-positive decode plan before allocating the output
buffer, annotate _decode_chunks as an Iterator, and document the
encode_pixels CPU dtype/range contract.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-21 02:38:35 +00:00
H1yori233 74b409d7cf [test]: add H3 VAE parity and memory benchmark 2026-08-19 02:45:23 -07:00
H1yori233 528cef02c4 optimize VAE memory 2026-08-12 02:48:43 -07:00
10 changed files with 899 additions and 39 deletions
+171 -25
View File
@@ -7,6 +7,7 @@ This module intentionally uses only PyTorch and FastVideo configuration types.
"""
import math
from collections.abc import Iterator
from dataclasses import dataclass
import torch
@@ -747,43 +748,140 @@ class AutoencoderKLMiniMaxH3(nn.Module):
moments = moments[:, :, :-self.config.token_drop]
return moments
def _decode(self, z: torch.Tensor) -> torch.Tensor:
tokens_chunk_size = self.tokens_chunk_size
def _encode_pixels(self, pixels: torch.Tensor) -> torch.Tensor:
"""Encode unnormalized pixels while keeping full videos off the accelerator."""
clip_length = self.config.clip_length
moments = []
for frame_start in range(0, pixels.shape[2], clip_length):
clip = pixels[:, :, frame_start:frame_start + clip_length].to(
device=self.pixel_mean.device,
dtype=torch.float32,
)
if pixels.dtype == torch.uint8:
clip = clip / 255.0
if clip.shape[2] < clip_length:
pad_frames = clip[:, :, -1:].repeat(1, 1, clip_length - clip.shape[2], 1, 1)
clip = torch.cat([clip, pad_frames], dim=2)
clip = self.normalize_pixels(clip)
moments.append(self._encode_clip(clip))
del clip
encoded = torch.cat(moments, dim=2)
if self.config.token_drop > 0:
encoded = encoded[:, :, :-self.config.token_drop]
return encoded
def _temporal_decode_plan(self, latent_num_frames: int) -> tuple[int, int, int]:
"""Return pad tokens, chunk count, and exact decoded frame count."""
if latent_num_frames <= 0:
raise ValueError(f"MiniMax-H3 latent frame count must be positive, got {latent_num_frames}.")
token_drop = self.config.token_drop
tokens_chunk_size = self.tokens_chunk_size
temporal_ratio = self.temporal_compression_ratio
chunk_num_frames = tokens_chunk_size * temporal_ratio
num_tokens = z.shape[2] + token_drop
num_tokens = latent_num_frames + token_drop
pad_tokens = (-num_tokens) % tokens_chunk_size
num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
if num_chunks < 1:
pad_tokens += tokens_chunk_size
num_chunks = 1
decoded_num_frames = num_chunks * (tokens_chunk_size * temporal_ratio - self.frame_pre_padding)
if token_drop > 0:
decoded_num_frames += self.frame_overlap
if pad_tokens > 0:
intra_tail = self.config.clip_length % temporal_ratio
pad_frames = sum(intra_tail if intra_tail and (latent_num_frames + offset) %
tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
decoded_num_frames -= pad_frames
if decoded_num_frames <= 0:
raise RuntimeError(
f"MiniMax-H3 decode plan produced {decoded_num_frames} frames for {latent_num_frames} latent "
"frames; the clip_length/token_drop configuration is inconsistent.")
return pad_tokens, num_chunks, decoded_num_frames
def _decode_chunks(self, z: torch.Tensor) -> Iterator[torch.Tensor]:
"""Yield finalized temporal chunks in decode order."""
tokens_chunk_size = self.tokens_chunk_size
chunk_num_frames = tokens_chunk_size * self.temporal_compression_ratio
pad_tokens, num_chunks, output_num_frames = self._temporal_decode_plan(z.shape[2])
if pad_tokens > 0:
z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
decoded_chunks = []
output_frame_start = 0
overlap = None
for index in range(num_chunks):
start = index * tokens_chunk_size
clip = self._decode_clip(z[:, :, start:start + tokens_chunk_size + self.token_overlap])
for overlap_index in range(int(token_drop > 0) + 1):
frame_start = overlap_index * chunk_num_frames
chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
chunk = chunk[:, :, self.frame_pre_padding:]
if overlap_index == 0:
if overlap is not None:
chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
decoded_chunks.append(chunk)
else:
overlap = chunk
if overlap is not None:
decoded_chunks.append(overlap)
decoded = torch.cat(decoded_chunks, dim=2)
chunk = clip[:, :, self.frame_pre_padding:chunk_num_frames]
next_overlap = None
if self.config.token_drop > 0:
next_overlap = clip[:, :, chunk_num_frames + self.frame_pre_padding:].clone()
if overlap is not None:
chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
if pad_tokens > 0:
intra_tail = self.config.clip_length % temporal_ratio
num_tokens_before_pad = z.shape[2] - pad_tokens
pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
decoded = decoded[:, :, :-pad_frames]
return decoded
num_frames = min(chunk.shape[2], output_num_frames - output_frame_start)
if num_frames > 0:
yield chunk[:, :, :num_frames]
output_frame_start += num_frames
overlap = next_overlap
if overlap is not None and output_frame_start < output_num_frames:
yield overlap[:, :, :output_num_frames - output_frame_start]
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,
)
@staticmethod
def _streams_chunk_copies(z: torch.Tensor, output: torch.Tensor) -> bool:
"""Whether finalized chunks copy to ``output`` asynchronously on the current CUDA stream."""
return z.device.type == "cuda" and output.is_pinned()
def _decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> None:
"""Decode temporal chunks and immediately copy finalized pixels to CPU.
Device-to-host copies run per (batch, channel) plane: the temporal
slice of ``output`` is strided across channels, but each plane is
contiguous on both sides, so every transfer stays a direct memcpy
instead of staging through a pageable CPU temporary. With a pinned
``output`` the copies are additionally asynchronous and overlap the
next chunk's decode; ``decode_to_pixels`` synchronizes once before
returning.
"""
non_blocking = self._streams_chunk_copies(z, output)
output_frame_start = 0
for chunk in self._decode_chunks(z):
num_frames = chunk.shape[2]
pixels = self.denormalize_pixels(chunk.float()).clamp_(0, 1)
target = output[:, :, output_frame_start:output_frame_start + num_frames]
if z.device.type == "cuda":
pixels = pixels.contiguous()
for batch_index in range(pixels.shape[0]):
for channel_index in range(pixels.shape[1]):
target[batch_index, channel_index].copy_(pixels[batch_index, channel_index],
non_blocking=non_blocking)
else:
target.copy_(pixels)
output_frame_start += num_frames
if output_frame_start != output.shape[2]:
raise RuntimeError(
f"MiniMax-H3 decode wrote {output_frame_start} frames into an output buffer expecting "
f"{output.shape[2]}.")
def _decode(self, z: torch.Tensor) -> torch.Tensor:
return torch.cat(list(self._decode_chunks(z)), dim=2)
def encode(
self,
@@ -799,6 +897,34 @@ class AutoencoderKLMiniMaxH3(nn.Module):
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 encode_keyframe(
self,
x: torch.Tensor,
@@ -825,6 +951,26 @@ class AutoencoderKLMiniMaxH3(nn.Module):
return (decoded, )
return DecoderOutput(sample=decoded)
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 forward(
self,
sample: torch.Tensor,
@@ -7,7 +7,7 @@ from typing import Any
import torch
from fastvideo.distributed import get_local_torch_device
from fastvideo.distributed import get_local_torch_device, get_world_group, model_parallel_is_initialized
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE
from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
@@ -21,6 +21,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import StageValidators as V
from fastvideo.pipelines.stages.validators import VerificationResult
from fastvideo.utils import is_pin_memory_available
def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
@@ -54,6 +55,13 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if model_parallel_is_initialized() and not get_world_group().is_first_rank:
# Distributed executors consume rank 0's ForwardBatch. Keep a
# verifier-compatible placeholder on other ranks and avoid
# duplicating the full VAE decode and CPU output buffer.
batch.output = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
return batch
layout = _layout(batch)
if batch.latents is None or batch.raw_latent_shape is None or len(batch.raw_latent_shape) != 5:
raise ValueError("MiniMax-H3 video latents or raw geometry are missing at decode.")
@@ -74,10 +82,16 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
batch.output = latents.detach().float().cpu()
return batch
output = torch.empty(
self.vae.decoded_pixel_shape(latents.shape),
device="cpu",
dtype=torch.float32,
pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
)
# The published decode recipe uses FP16 autocast over FP32 weights.
with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
video = self.vae.decode(latents).sample
batch.output = self.vae.denormalize_pixels(video.float()).clamp_(0, 1).cpu()
self.vae.decode_to_pixels(latents, output)
batch.output = output
return batch
finally:
if fastvideo_args.vae_cpu_offload:
@@ -107,6 +121,12 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if model_parallel_is_initialized() and not get_world_group().is_first_rank:
batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
self._clear_runtime(batch)
return batch
layout = _layout(batch)
if batch.audio_latents is None:
raise ValueError("MiniMax-H3 audio latents are missing at decode.")
@@ -119,9 +119,8 @@ class MiniMaxH3LatentPreparationStage(PipelineStage):
if reference.frames is None:
raise ValueError("MiniMax-H3 reference video frames are missing.")
frames = reference.frames[:trim_reference_num_frames(reference.frames.shape[0])]
pixels = torch.from_numpy(frames.copy()).permute(3, 0, 1, 2)[None]
pixels = pixels.to(device=device, dtype=torch.float32).div_(255.0)
posterior = self.vae.encode(self.vae.normalize_pixels(pixels)).latent_dist
pixels = torch.from_numpy(np.ascontiguousarray(frames)).permute(3, 0, 1, 2)[None]
posterior = self.vae.encode_pixels(pixels).latent_dist
latents = self.vae.normalize_latents(_sample_visual_posterior(posterior).to(
torch.float16).float()).cpu()
reference.num_latent_frames = int(latents.shape[2])
@@ -0,0 +1,130 @@
# SPDX-License-Identifier: Apache-2.0
from types import SimpleNamespace
import numpy as np
import torch
from fastvideo.pipelines.basic.minimax_h3.packing import (
MiniMaxH3PackedLayout,
patchify_video_latents,
)
from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_decoding
from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3VideoDecodingStage
from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import (
MINIMAX_H3_LAYOUT_KEY,
MiniMaxH3LatentPreparationStage,
)
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
def _layout(rows: int, latent_shape: tuple[int, ...]) -> MiniMaxH3PackedLayout:
empty = torch.empty(0, dtype=torch.long)
return MiniMaxH3PackedLayout(
sequence_length=rows,
position_ids=empty,
token_tags=empty,
video_indices=empty,
audio_indices=empty,
text_indices=empty,
num_condition_video_rows=0,
num_condition_audio_rows=0,
num_video_latent_frames=latent_shape[2],
latent_height=latent_shape[3],
latent_width=latent_shape[4],
num_audio_latents=0,
)
def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
observed = {}
class VAE:
def encode_pixels(self, pixels):
observed["pixels"] = pixels
posterior = SimpleNamespace(sample=lambda generator=None: torch.zeros(1, 4, 7, 4, 4))
return SimpleNamespace(latent_dist=posterior)
def normalize_latents(self, latents):
return latents
stage = MiniMaxH3LatentPreparationStage(
transformer=SimpleNamespace(patch_size=(1, 1, 1)),
vae=VAE(),
audio_vae=None,
scheduler=None,
ref2va=True,
)
reference = MiniMaxH3PreparedReference(
media_type="video",
frames=np.zeros((22, 16, 16, 3), dtype=np.uint8),
)
rows = stage._encode_visual_rows([reference], torch.device("cpu"))
assert observed["pixels"].dtype == torch.uint8
assert observed["pixels"].device.type == "cpu"
assert rows[0].shape == (7 * 4 * 4, 4)
def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
latent_shape = (1, 4, 2, 4, 4)
latents = torch.randn(latent_shape)
rows = patchify_video_latents(latents, (1, 1, 1))
batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape)
observed = {}
class VAE:
def to(self, device):
return self
def denormalize_latents(self, decoded_latents):
return decoded_latents
def decoded_pixel_shape(self, shape):
assert tuple(shape) == latent_shape
return (1, 3, 5, 16, 16)
def decode_to_pixels(self, decoded_latents, output):
observed["latents"] = decoded_latents
observed["output"] = output
output.fill_(0.25)
monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(
batch,
SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False),
)
torch.testing.assert_close(observed["latents"], latents)
assert observed["output"] is result.output
assert result.output.device.type == "cpu"
assert torch.all(result.output == 0.25)
def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
class VAE:
sampling_rate = 32000
def to(self, device):
raise AssertionError("non-output ranks must not execute a VAE")
monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True)
monkeypatch.setattr(minimax_h3_decoding, "get_world_group", lambda: SimpleNamespace(is_first_rank=False))
args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True)
video = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace()).forward(ForwardBatch(data_type="video"), args)
assert video.output.shape == (0, 3, 0, 0, 0)
audio_batch = ForwardBatch(data_type="audio", latents=torch.zeros(1), audio_latents=torch.zeros(1))
audio_batch.extra[MINIMAX_H3_LAYOUT_KEY] = object()
audio = minimax_h3_decoding.MiniMaxH3AudioDecodingStage(VAE()).forward(audio_batch, args)
assert audio.extra["audio"].shape == (0, 2)
assert audio.extra["audio_sample_rate"] == 32000
assert audio.latents is None
assert audio.audio_latents is None
assert MINIMAX_H3_LAYOUT_KEY not in audio.extra
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
from torch.testing import assert_close
from fastvideo.configs.models.vaes.minimax_h3_video import (
MiniMaxH3VideoVAEArchConfig,
MiniMaxH3VideoVAEConfig,
)
from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
def _tiny_vae() -> AutoencoderKLMiniMaxH3:
arch = MiniMaxH3VideoVAEArchConfig(
latent_channels=4,
block_out_channels=(32, 32),
layers_per_block=1,
spatial_downsample_factors=(2, 2),
temporal_downsample_factors=(2, 2),
decoder_num_layers=1,
decoder_num_attention_heads=1,
decoder_attention_head_dim=8,
decoder_num_register_tokens=2,
decoder_ffn_mult=1,
latents_mean=(0.0, ) * 4,
latents_std=(1.0, ) * 4,
)
return AutoencoderKLMiniMaxH3(
MiniMaxH3VideoVAEConfig(
arch_config=arch,
use_tiling=False,
use_temporal_tiling=False,
use_parallel_tiling=False,
)).eval()
@torch.inference_mode()
def test_encode_pixels_matches_encode() -> None:
torch.manual_seed(20260810)
vae = _tiny_vae()
pixels = torch.randint(0, 256, (1, 3, 22, 16, 16), dtype=torch.uint8)
expected = vae.encode(vae.normalize_pixels(pixels.float().div(255))).latent_dist.parameters
assert_close(vae.encode_pixels(pixels).latent_dist.parameters, expected, atol=0.0, rtol=0.0)
float_pixels = torch.rand(1, 3, 22, 16, 16)
original_pixels = float_pixels.clone()
expected = vae.encode(vae.normalize_pixels(float_pixels)).latent_dist.parameters
actual = vae.encode_pixels(float_pixels).latent_dist.parameters
assert_close(float_pixels, original_pixels, atol=0.0, rtol=0.0)
assert_close(actual, expected, atol=0.0, rtol=0.0)
with pytest.raises(ValueError, match="must remain on CPU"):
vae.encode_pixels(torch.empty(1, 3, 1, 16, 16, device="meta"))
def _legacy_decode(vae: AutoencoderKLMiniMaxH3, z: torch.Tensor) -> torch.Tensor:
"""Verbatim pre-streaming ``_decode`` (main @ 8208536cd) as a reference oracle."""
tokens_chunk_size = vae.tokens_chunk_size
token_drop = vae.config.token_drop
temporal_ratio = vae.temporal_compression_ratio
chunk_num_frames = tokens_chunk_size * temporal_ratio
num_tokens = z.shape[2] + token_drop
pad_tokens = (-num_tokens) % tokens_chunk_size
num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
if pad_tokens > 0:
z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
decoded_chunks = []
overlap = None
for index in range(num_chunks):
start = index * tokens_chunk_size
clip = vae._decode_clip(z[:, :, start:start + tokens_chunk_size + vae.token_overlap])
for overlap_index in range(int(token_drop > 0) + 1):
frame_start = overlap_index * chunk_num_frames
chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
chunk = chunk[:, :, vae.frame_pre_padding:]
if overlap_index == 0:
if overlap is not None:
chunk = vae._blend(overlap, chunk, vae.frame_overlap, dim=-3)
decoded_chunks.append(chunk)
else:
overlap = chunk
if overlap is not None:
decoded_chunks.append(overlap)
decoded = torch.cat(decoded_chunks, dim=2)
if pad_tokens > 0:
intra_tail = vae.config.clip_length % temporal_ratio
num_tokens_before_pad = z.shape[2] - pad_tokens
pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
decoded = decoded[:, :, :-pad_frames]
return decoded
@pytest.mark.parametrize("latent_frames", (2, 12))
@torch.inference_mode()
def test_decode_to_pixels_matches_decode(latent_frames: int) -> None:
torch.manual_seed(20260810)
vae = _tiny_vae()
latents = torch.randn(1, 4, latent_frames, 4, 4)
expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1)
actual = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
vae.decode_to_pixels(latents, actual)
assert_close(actual, expected, atol=0.0, rtol=0.0)
# 3: one chunk with pad tokens; 6: pad hitting the intra-clip tail;
# 12: two blended chunks without padding; 13: three chunks plus pad trim.
@pytest.mark.parametrize("latent_frames", (3, 6, 12, 13))
@torch.inference_mode()
def test_decode_matches_legacy_algorithm(latent_frames: int) -> None:
"""The chunk iterator must stay bit-exact with the pre-streaming decode."""
torch.manual_seed(20260811 + latent_frames)
vae = _tiny_vae()
latents = torch.randn(1, 4, latent_frames, 4, 4)
expected = _legacy_decode(vae, latents)
decoded = vae.decode(latents).sample
assert_close(decoded, expected, atol=0.0, rtol=0.0)
streamed = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
vae.decode_to_pixels(latents, streamed)
assert_close(streamed, vae.denormalize_pixels(expected.float()).clamp_(0, 1), atol=0.0, rtol=0.0)
@torch.inference_mode()
def test_streaming_slicing_matches_unbatched() -> None:
torch.manual_seed(20260812)
vae = _tiny_vae()
vae.enable_slicing()
latents = torch.randn(2, 4, 7, 4, 4)
expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1)
actual = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
vae.decode_to_pixels(latents, actual)
assert_close(actual, expected, atol=0.0, rtol=0.0)
pixels = torch.randint(0, 256, (2, 3, 22, 16, 16), dtype=torch.uint8)
expected_moments = vae.encode(vae.normalize_pixels(pixels.float().div(255))).latent_dist.parameters
assert_close(vae.encode_pixels(pixels).latent_dist.parameters, expected_moments, atol=0.0, rtol=0.0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="pinned-buffer streaming requires CUDA")
@torch.inference_mode()
def test_decode_to_pixels_pinned_buffer_matches_dense_on_cuda() -> None:
"""Async chunk copies into a pinned buffer must equal the dense decode."""
torch.manual_seed(20260813)
vae = _tiny_vae().to("cuda")
latents = torch.randn(1, 4, 12, 4, 4, device="cuda")
expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1).cpu()
pinned = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32, pin_memory=True)
vae.decode_to_pixels(latents, pinned)
assert_close(pinned, expected, atol=0.0, rtol=0.0)
def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None:
vae = _tiny_vae()
latents = torch.randn(1, 4, 2, 4, 4)
output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
monkeypatch.setattr(vae, "_decode_chunks", lambda _: iter(()))
with pytest.raises(RuntimeError, match="wrote 0 frames"):
vae.decode_to_pixels(latents, output)
@@ -14,6 +14,7 @@
| Qwen3-VL encoder | exact text/image/video hidden states through the production loader | complete |
| FL2VA and Ref2VA DiTs | exact video/audio heads for both model partitions | complete |
| Video VAE | exact encode, normalization, and decode through the production loader | complete |
| Video VAE streaming | exact chunked encode/decode and output-rank-only distributed decode | complete |
| Audio VAE | exact encode and normalization; decode maximum absolute drift `2.4e-7` | complete |
| Video/audio schedulers | pinned `12/3` schedule parity | complete |
| FL2VA packing | pinned row, position, tag, timestep, and RNG parity | complete |
@@ -34,6 +35,7 @@ T2VA, FL2VA, and Ref2VA match the official video/audio latents exactly.
- Load `transformer/` for T2VA/FL2VA and `transformer_ref/` for Ref2VA.
- Keep `last_image`, `references`, and `audio_latents` on the typed request path.
- Treat the published component folders as the loading boundary.
- Keep reference videos on CPU between VAE clips and decode final pixels only on the executor's output rank.
## Evidence boundary
+33
View File
@@ -12,6 +12,14 @@ FastVideo-owned unit contracts belong under `fastvideo/tests/`.
The reference helper verifies the pinned source and import origin. A missing checkout may skip a source-parity module;
that skip is not parity evidence.
## FastVideo unit contracts
```bash
pytest \
fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py \
fastvideo/tests/stages/test_minimax_h3_vae_streaming.py -q
```
## Registry smoke
```bash
@@ -51,5 +59,30 @@ pytest \
With a gate enabled, missing CUDA, source, or weights is a failure. Recorded component evidence is exact for both DiT
partitions, the video VAE, and all Qwen3-VL hidden states; audio decode has maximum absolute drift `2.4e-7`.
The video VAE test verifies the reference checkout at commit
`abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc` and compares the production CPU `uint8` `encode_pixels()` path against
the official posterior element by element.
## Video VAE memory benchmark
The benchmark uses one warmup and three measured runs with `vae_cpu_offload=True`. It reports absolute and
stage-incremental allocated/reserved CUDA peaks for every rank. For SP runs, the reported aggregate is explicitly the
sum of rank-local maxima, not a simultaneous node peak.
```bash
python tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \
--source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \
--revision-label candidate --operation encode
python -m torch.distributed.run --nproc_per_node=4 \
tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \
--source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \
--revision-label candidate-sp4 --operation decode
```
Run the same script with `--source-root` pointed at the base checkout for a comparable baseline. The default workload
is deterministic `124 x 768 x 1344` video geometry with seed `20260803`; the JSON record includes source/model
revisions, software/allocator metadata, exact measurement boundaries, per-repetition values, and output shapes.
FastVideo joint audio/video generation and SP=1/SP=4 latent consistency have been validated. T2VA, FL2VA, and
Ref2VA video/audio latents match the pinned Diffusers pipeline exactly.
@@ -3,6 +3,7 @@
import hashlib
import inspect
import os
import subprocess
from pathlib import Path
import pytest
@@ -13,6 +14,28 @@ REFERENCE_SRC = REFERENCE_ROOT / "src"
PINNED_COMMIT = "abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc"
def _run_git(*args: str) -> str:
try:
result = subprocess.run(
["git", "-C", str(REFERENCE_ROOT), *args],
check=True,
capture_output=True,
text=True,
)
except (FileNotFoundError, subprocess.CalledProcessError) as error:
raise RuntimeError(f"Could not verify the MiniMax-H3 reference checkout at {REFERENCE_ROOT}.") from error
return result.stdout.strip()
def assert_reference_revision() -> None:
actual_commit = _run_git("rev-parse", "HEAD")
if actual_commit != PINNED_COMMIT:
raise RuntimeError(f"MiniMax-H3 parity requires Diffusers commit {PINNED_COMMIT}, got {actual_commit}.")
dirty_source = _run_git("status", "--short", "--untracked-files=no", "--", "src/diffusers")
if dirty_source:
raise RuntimeError(f"MiniMax-H3 reference source has tracked changes:\n{dirty_source}")
def assert_pinned_reference(relative_path: str, sha256: str) -> Path:
path = REFERENCE_ROOT / relative_path
if not path.is_file():
@@ -0,0 +1,300 @@
# SPDX-License-Identifier: Apache-2.0
"""Measure MiniMax-H3 production VAE stage memory with CPU offload enabled."""
from __future__ import annotations
import argparse
import gc
import hashlib
import json
import os
from pathlib import Path
import statistics
import subprocess
import sys
import time
from types import SimpleNamespace
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--source-root", type=Path, required=True)
parser.add_argument("--model-root", type=Path, required=True)
parser.add_argument("--revision-label", required=True)
parser.add_argument("--operation", choices=("encode", "decode"), required=True)
parser.add_argument("--height", type=int, default=768)
parser.add_argument("--width", type=int, default=1344)
parser.add_argument("--num-frames", type=int, default=124)
parser.add_argument("--seed", type=int, default=20260803)
parser.add_argument("--warmups", type=int, default=1)
parser.add_argument("--repetitions", type=int, default=3)
return parser.parse_args()
def _git(source_root: Path, *args: str) -> str:
result = subprocess.run(
["git", "-C", str(source_root), *args],
check=True,
capture_output=True,
text=True,
)
return result.stdout.strip()
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as file:
for block in iter(lambda: file.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def _model_snapshot(model_root: Path) -> str | None:
parts = model_root.resolve().parts
try:
return parts[parts.index("snapshots") + 1]
except (ValueError, IndexError):
return None
def _make_layout(rows, latent_shape):
import torch
from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout
empty = torch.empty(0, dtype=torch.long, device=rows.device)
return MiniMaxH3PackedLayout(
sequence_length=rows.shape[0],
position_ids=empty,
token_tags=empty,
video_indices=empty,
audio_indices=empty,
text_indices=empty,
num_condition_video_rows=0,
num_condition_audio_rows=0,
num_video_latent_frames=latent_shape[2],
latent_height=latent_shape[3],
latent_width=latent_shape[4],
num_audio_latents=0,
)
def _build_operation(args, vae, device):
import numpy as np
import torch
from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
from fastvideo.pipelines.basic.minimax_h3.packing import patchify_video_latents, video_latent_num_frames
from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3VideoDecodingStage
from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import (
MINIMAX_H3_LAYOUT_KEY,
MiniMaxH3LatentPreparationStage,
)
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
patch_size = MiniMaxH3Config().arch_config.patch_size
transformer = SimpleNamespace(patch_size=patch_size)
runtime_args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True)
if args.operation == "encode":
if int(os.environ.get("WORLD_SIZE", "1")) != 1:
raise ValueError("The encode benchmark is single-rank; use one process.")
frames = np.random.default_rng(args.seed).integers(
0,
256,
size=(args.num_frames, args.height, args.width, 3),
dtype=np.uint8,
)
stage = MiniMaxH3LatentPreparationStage(
transformer=transformer,
vae=vae,
audio_vae=None,
scheduler=None,
ref2va=True,
)
def run_once():
reference = MiniMaxH3PreparedReference(media_type="video", frames=frames)
vae.to(device)
try:
return stage._encode_visual_rows([reference], device)[0]
finally:
vae.to("cpu")
return run_once, {
"input": "NumPy PCG64 CPU uint8 RGB pixels",
"input_shape": [1, 3, args.num_frames, args.height, args.width],
"boundary": "before VAE CPU-to-GPU transfer through post-encode VAE CPU offload",
}
latent_shape = (
1,
vae.latent_channels,
video_latent_num_frames(args.num_frames),
args.height // vae.spatial_compression_ratio,
args.width // vae.spatial_compression_ratio,
)
generator = torch.Generator(device=device).manual_seed(args.seed)
latents = torch.randn(latent_shape, generator=generator, device=device, dtype=torch.float32)
rows = patchify_video_latents(latents, patch_size)
layout = _make_layout(rows, latent_shape)
stage = MiniMaxH3VideoDecodingStage(vae, transformer)
def run_once():
batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout
return stage.forward(batch, runtime_args).output
return run_once, {
"input": "PyTorch Philox normalized FP32 latents",
"input_shape": list(latent_shape),
"boundary": "MiniMaxH3VideoDecodingStage.forward including VAE CPU-to-GPU transfer and CPU offload",
}
def _measure(run_once, device, warmups: int, repetitions: int) -> dict:
import torch
for _ in range(warmups):
with torch.inference_mode():
result = run_once()
torch.cuda.synchronize(device)
del result
gc.collect()
torch.cuda.empty_cache()
records = []
for repetition in range(repetitions):
gc.collect()
torch.cuda.empty_cache()
torch.cuda.synchronize(device)
start_allocated = torch.cuda.memory_allocated(device)
start_reserved = torch.cuda.memory_reserved(device)
torch.cuda.reset_peak_memory_stats(device)
started = time.perf_counter()
with torch.inference_mode():
result = run_once()
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - started
records.append({
"repetition": repetition,
"elapsed_seconds": elapsed,
"start_allocated_bytes": start_allocated,
"peak_allocated_bytes": torch.cuda.max_memory_allocated(device),
"incremental_peak_allocated_bytes": torch.cuda.max_memory_allocated(device) - start_allocated,
"start_reserved_bytes": start_reserved,
"peak_reserved_bytes": torch.cuda.max_memory_reserved(device),
"incremental_peak_reserved_bytes": torch.cuda.max_memory_reserved(device) - start_reserved,
"output_shape": list(result.shape),
})
del result
return {
"repetitions": records,
"median_elapsed_seconds": statistics.median(record["elapsed_seconds"] for record in records),
"max_peak_allocated_bytes": max(record["peak_allocated_bytes"] for record in records),
"max_incremental_peak_allocated_bytes": max(
record["incremental_peak_allocated_bytes"] for record in records),
"max_peak_reserved_bytes": max(record["peak_reserved_bytes"] for record in records),
"max_incremental_peak_reserved_bytes": max(
record["incremental_peak_reserved_bytes"] for record in records),
}
def main() -> None:
args = _parse_args()
if args.warmups < 0 or args.repetitions < 1:
raise ValueError("warmups must be non-negative and repetitions must be positive.")
source_root = args.source_root.resolve()
sys.path.insert(0, str(source_root))
import torch
import torch.distributed as dist
import fastvideo
from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
from fastvideo.models.loader.component_loader import VAELoader
imported_root = Path(fastvideo.__file__).resolve().parents[1]
if imported_root != source_root:
raise RuntimeError(f"Imported FastVideo from {imported_root}, expected {source_root}.")
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29673")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")
world_size = int(os.environ["WORLD_SIZE"])
maybe_init_distributed_environment_and_model_parallel(1, world_size)
rank = dist.get_rank()
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
component_dir = args.model_root.resolve() / "vae"
config_path = component_dir / "config.json"
index_path = component_dir / "diffusion_pytorch_model.safetensors.index.json"
for path in (config_path, index_path):
if not path.is_file():
raise FileNotFoundError(path)
loader_args = SimpleNamespace(
pipeline_config=MiniMaxH3PipelineConfig(),
model_paths={},
vae_cpu_offload=True,
)
vae = VAELoader().load(str(component_dir), loader_args)
parameter_bytes = sum(parameter.numel() * parameter.element_size() for parameter in vae.parameters())
run_once, workload = _build_operation(args, vae, device)
rank_result = _measure(run_once, device, args.warmups, args.repetitions)
rank_result.update({
"rank": rank,
"device_index": device.index,
"device_name": torch.cuda.get_device_name(device),
"device_total_memory_bytes": torch.cuda.get_device_properties(device).total_memory,
})
rank_results = [None] * world_size
dist.all_gather_object(rank_results, rank_result)
if rank == 0:
result = {
"schema_version": 1,
"exit_status": 0,
"revision_label": args.revision_label,
"source_root": str(source_root),
"source_git_head": _git(source_root, "rev-parse", "HEAD"),
"source_tracked_status": _git(source_root, "status", "--short", "--untracked-files=no"),
"benchmark_script_sha256": _sha256(Path(__file__).resolve()),
"model_root": str(args.model_root.resolve()),
"model_snapshot": _model_snapshot(args.model_root),
"model_config_sha256": _sha256(config_path),
"model_index_sha256": _sha256(index_path),
"vae_parameter_bytes": parameter_bytes,
"operation": args.operation,
"vae_cpu_offload": True,
"pin_cpu_memory": False,
"warmups": args.warmups,
"measurement_repetitions": args.repetitions,
"seed": args.seed,
"workload": workload,
"world_size": world_size,
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"allocator_config": os.environ.get("PYTORCH_ALLOC_CONF")
or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"),
"rank_results": rank_results,
"sum_rank_local_max_peak_allocated_bytes": sum(
item["max_peak_allocated_bytes"] for item in rank_results),
"sum_rank_local_max_incremental_peak_allocated_bytes": sum(
item["max_incremental_peak_allocated_bytes"] for item in rank_results),
"sum_rank_local_max_peak_reserved_bytes": sum(
item["max_peak_reserved_bytes"] for item in rank_results),
"sum_rank_local_max_incremental_peak_reserved_bytes": sum(
item["max_incremental_peak_reserved_bytes"] for item in rank_results),
}
print("MINIMAX_H3_VAE_MEMORY=" + json.dumps(result, sort_keys=True), flush=True)
cleanup_dist_env_and_memory()
if __name__ == "__main__":
main()
@@ -62,8 +62,9 @@ def _require_assets() -> tuple[torch.device, Path]:
def _load_official(component_dir: Path, device: torch.device) -> torch.nn.Module:
from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import AutoencoderKLMiniMaxH3
from tests.local_tests.minimax_h3._reference import assert_reference_source
from tests.local_tests.minimax_h3._reference import assert_reference_revision, assert_reference_source
assert_reference_revision()
assert_reference_source(
AutoencoderKLMiniMaxH3,
"src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py",
@@ -101,14 +102,35 @@ def _load_fastvideo(component_dir: Path) -> torch.nn.Module:
return model
def _make_video() -> torch.Tensor:
def _make_pixels() -> torch.Tensor:
generator = torch.Generator(device="cpu").manual_seed(20260803)
# Two logical 17-frame clips after padding are required for H3's three-token
# temporal overlap contract; 32 px remains safely above reflect-pad minima.
return torch.randn(1, 3, 22, 32, 32, generator=generator, dtype=torch.float32)
return torch.randint(0, 256, (1, 3, 22, 32, 32), generator=generator, dtype=torch.uint8)
def _run(model: torch.nn.Module, video: torch.Tensor) -> dict[str, torch.Tensor]:
def _normalize_pixels(pixels: torch.Tensor, device: torch.device) -> torch.Tensor:
video = pixels.to(device=device, dtype=torch.float32).div_(255.0)
pixel_mean = torch.tensor((0.485, 0.456, 0.406), device=device).view(1, -1, 1, 1, 1)
pixel_std = torch.tensor((0.229, 0.224, 0.225), device=device).view(1, -1, 1, 1, 1)
return (video - pixel_mean) / pixel_std
def _run_encode_pixels(model: torch.nn.Module, pixels: torch.Tensor) -> dict[str, torch.Tensor]:
with torch.inference_mode():
posterior = model.encode_pixels(pixels, return_dict=False)[0]
return {
"mean": posterior.mean.detach().cpu(),
"logvar": posterior.logvar.detach().cpu(),
}
def _run(
model: torch.nn.Module,
video: torch.Tensor,
*,
stream_output: bool = False,
) -> dict[str, torch.Tensor]:
with torch.inference_mode():
posterior = model.encode(video, return_dict=False)[0]
latents = posterior.mode()
@@ -121,12 +143,25 @@ def _run(model: torch.nn.Module, video: torch.Tensor) -> dict[str, torch.Tensor]
dtype=latents.dtype).view(1, -1, 1, 1, 1)
normalized = (latents - mean) / std
decoded = model.decode(latents, return_dict=False)[0]
return {
pixel_mean = torch.tensor((0.485, 0.456, 0.406), device=decoded.device,
dtype=torch.float32).view(1, -1, 1, 1, 1)
pixel_std = torch.tensor((0.229, 0.224, 0.225), device=decoded.device,
dtype=torch.float32).view(1, -1, 1, 1, 1)
pixels = (decoded.float() * pixel_std + pixel_mean).clamp_(0, 1)
streamed = None
if stream_output:
streamed = torch.empty(model.decoded_pixel_shape(latents.shape), dtype=torch.float32, device="cpu")
model.decode_to_pixels(latents, streamed)
result = {
"mean": posterior.mean.detach().cpu(),
"logvar": posterior.logvar.detach().cpu(),
"normalized": normalized.detach().cpu(),
"decoded": decoded.detach().cpu(),
"pixels": pixels.detach().cpu(),
}
if streamed is not None:
result["streamed"] = streamed
return result
def _reclaim_vram() -> None:
@@ -151,15 +186,16 @@ def _assert_tensor_parity(name: str, actual: torch.Tensor, expected: torch.Tenso
def test_minimax_h3_video_vae_parity() -> None:
"""Match posterior, normalization, geometry, and deterministic decode."""
device, component_dir = _require_assets()
video = _make_video()
pixels = _make_pixels()
official = _load_official(component_dir, device)
expected = _run(official, video.to(device))
expected = _run(official, _normalize_pixels(pixels, device))
del official
_reclaim_vram()
fastvideo = _load_fastvideo(component_dir)
actual = _run(fastvideo, video.to(device))
actual = _run(fastvideo, _normalize_pixels(pixels, device), stream_output=True)
streaming_encode = _run_encode_pixels(fastvideo, pixels)
assert fastvideo.temporal_compression_ratio == 4
assert fastvideo.spatial_compression_ratio == 16
del fastvideo
@@ -167,5 +203,8 @@ def test_minimax_h3_video_vae_parity() -> None:
_assert_tensor_parity("video_vae.mean", actual["mean"], expected["mean"], 0.0)
_assert_tensor_parity("video_vae.logvar", actual["logvar"], expected["logvar"], 0.0)
_assert_tensor_parity("video_vae.streaming_encode.mean", streaming_encode["mean"], expected["mean"], 0.0)
_assert_tensor_parity("video_vae.streaming_encode.logvar", streaming_encode["logvar"], expected["logvar"], 0.0)
_assert_tensor_parity("video_vae.normalized", actual["normalized"], expected["normalized"], 0.0)
_assert_tensor_parity("video_vae.decode", actual["decoded"], expected["decoded"], 0.0)
_assert_tensor_parity("video_vae.streaming_decode", actual["streamed"], expected["pixels"], 0.0)