diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index d4c6b95a6..e2d671578 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -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, diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 5d77ac456..06bd1b88e 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -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.") diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index 34f4f9cba..40150fea8 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -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]) diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py new file mode 100644 index 000000000..1eab19ed4 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -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 diff --git a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py new file mode 100644 index 000000000..5a2c8783e --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -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) diff --git a/tests/local_tests/minimax_h3/PORT_STATUS.md b/tests/local_tests/minimax_h3/PORT_STATUS.md index 3a7fde89d..9c9ed2cee 100644 --- a/tests/local_tests/minimax_h3/PORT_STATUS.md +++ b/tests/local_tests/minimax_h3/PORT_STATUS.md @@ -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 diff --git a/tests/local_tests/minimax_h3/README.md b/tests/local_tests/minimax_h3/README.md index 0d1d78227..82f26c85f 100644 --- a/tests/local_tests/minimax_h3/README.md +++ b/tests/local_tests/minimax_h3/README.md @@ -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. diff --git a/tests/local_tests/minimax_h3/_reference.py b/tests/local_tests/minimax_h3/_reference.py index 2dbb77d21..ac80d7c84 100644 --- a/tests/local_tests/minimax_h3/_reference.py +++ b/tests/local_tests/minimax_h3/_reference.py @@ -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(): diff --git a/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py new file mode 100644 index 000000000..57e68a742 --- /dev/null +++ b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py @@ -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() diff --git a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py index 467528659..98308a441 100644 --- a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py +++ b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py @@ -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)