Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
942f7db3db | ||
|
|
aadb23f409 | ||
|
|
74b409d7cf | ||
|
|
528cef02c4 |
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user