diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 89b9b3953..3ebac451b 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -1041,10 +1041,16 @@ class VideoGenerator: # `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or # `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same # data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback. - vid_u8 = _quantize_video_frames_to_uint8(output_batch.output) - frames = [ - torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8 - ] + src_frames = output_batch.output + if (isinstance(src_frames, torch.Tensor) and src_frames.dim() == 5 and src_frames.dtype == torch.uint8 + and src_frames.shape[0] == 1 and src_frames.shape[2] > 1): + frames = list(src_frames[0].permute(1, 2, 3, 0).contiguous().cpu().numpy()) + else: + vid_u8 = _quantize_video_frames_to_uint8(src_frames) + frames = [ + torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() + for x in vid_u8 + ] postprocess_time = time.perf_counter() - postprocess_start logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time) if logging_info is not None: diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 56c9a8cb1..a1eabdcb4 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -753,6 +753,11 @@ FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW = EnvBool(False, category="performance", doc="Use transposed weight views in the MiniMax-H3 VAE INT8 " "projections.") +FASTVIDEO_H3_VAE_UINT8_PIXELS = EnvBool(True, + category="performance", + doc="Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per " + "chunk and stream uint8 into the host buffer; frames are bitwise identical " + "to the fp32 path. 0 streams fp32.") FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT = EnvBool(False, category="performance", doc="Fused dequantization epilogue for the MiniMax-H3 VAE INT8 " diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index fa1ecbdbf..224719d88 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -1016,7 +1016,7 @@ class AutoencoderKLMiniMaxH3(nn.Module): @staticmethod def _copy_chunk_pixels(pixels: torch.Tensor, output: torch.Tensor, frame_start: int, non_blocking: bool) -> None: - """Copy one finalized fp32 pixel chunk into the CPU ``output`` buffer. + """Copy one finalized fp32 pixel chunk into the CPU ``output`` buffer (fp32, or uint8 quantized on device). Device-to-host copies run per (batch, channel) plane: the temporal slice of ``output`` is strided across channels, but each plane is @@ -1027,6 +1027,9 @@ class AutoencoderKLMiniMaxH3(nn.Module): before releasing the buffer. """ target = output[:, :, frame_start:frame_start + pixels.shape[2]] + if output.dtype == torch.uint8: + # Same mul/clamp/truncating cast as the stage's fp32 host path, so frames are bitwise equal. + pixels = (pixels * 255).clamp_(0, 255).to(torch.uint8) if pixels.device.type == "cuda": pixels = pixels.contiguous() for batch_index in range(pixels.shape[0]): @@ -1145,11 +1148,17 @@ class AutoencoderKLMiniMaxH3(nn.Module): 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.""" + """Stream decoded ``[0, 1]`` pixels into a caller-owned CPU buffer. + + ``output`` is float32, or uint8 when the caller asks for the quantized + stream (``FASTVIDEO_H3_VAE_UINT8_PIXELS``); uint8 quantizes each chunk on + the GPU with the fp32 host path's exact mul/clamp/truncate arithmetic. + """ expected_shape = self.decoded_pixel_shape(z.shape) - if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape: + if (output.device.type != "cpu" or output.dtype not in (torch.float32, torch.uint8) + or tuple(output.shape) != expected_shape): raise ValueError( - "`output` must be a CPU float32 tensor with shape " + "`output` must be a CPU float32 or uint8 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: 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 fdfe8dd18..42e5acd30 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -7,6 +7,7 @@ from typing import Any import torch +import fastvideo.envs as envs from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger @@ -68,11 +69,16 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None: super().__init__() self.vae = vae - self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None + self._output_buffer: tuple[tuple[tuple[int, ...], bool, torch.dtype], torch.Tensor] | None = None self._unified_memory: bool | None = None - def _pixel_buffer(self, shape: tuple[int, ...], pin_memory: bool, device: torch.device) -> torch.Tensor: - """The CPU FP32 pixel buffer, reused across requests of the same geometry on unified memory. + def _pixel_buffer(self, + shape: tuple[int, ...], + pin_memory: bool, + device: torch.device, + dtype: torch.dtype = torch.float32) -> torch.Tensor: + """The CPU pixel buffer (FP32, or uint8 when chunks are quantized on the GPU), reused across requests of + the same geometry on unified memory. The stage returns a uint8 copy, so the buffer never escapes. Reuse matters on unified-memory GPUs (DGX Spark): device-to-host copies into never-touched pageable @@ -84,11 +90,11 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): if self._unified_memory is None: self._unified_memory = current_platform.has_unified_memory(device.index or 0) if not self._unified_memory: - return allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory) - key = (tuple(shape), bool(pin_memory)) + return allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory) + key = (tuple(shape), bool(pin_memory), dtype) if self._output_buffer is None or self._output_buffer[0] != key: self._output_buffer = None - buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory) + buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory) if not buffer.is_pinned(): buffer.zero_() self._output_buffer = (key, buffer) @@ -170,8 +176,9 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): output = None if is_output_rank: + pixel_dtype = torch.uint8 if envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.get() and not parallel else torch.float32 output = self._pixel_buffer(tuple(self.vae.decoded_pixel_shape(latents.shape)), - fastvideo_args.pin_cpu_memory, device) + fastvideo_args.pin_cpu_memory, device, pixel_dtype) # Attribute the streamed decoder computation while retaining # per-chunk device-to-host transfer and pinned-buffer reuse. # fp16 for the release decoders; NVFP4 decoders were trained under bf16. diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 7601989fd..f7c241083 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -75,6 +75,7 @@ def test_reference_video_encode_keeps_pixels_on_cpu() -> None: def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0") latent_shape = (1, 4, 2, 4, 4) latents = torch.randn(latent_shape) rows = patchify_video_latents(latents, (1, 1, 1)) @@ -129,6 +130,7 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: + monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0") latent_shape = (1, 4, 2, 4, 4) latents = torch.randn(latent_shape) rows = patchify_video_latents(latents, (1, 1, 1)) @@ -177,10 +179,50 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: } +def test_decode_stage_streams_uint8_by_default(monkeypatch) -> None: + latent_shape = (1, 4, 2, 4, 4) + rows = patchify_video_latents(torch.randn(latent_shape), (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): + return (1, 3, 5, 16, 16) + + def decode_to_pixels(self, decoded_latents, output): + observed["dtype"] = output.dtype + output.fill_(63) + + monkeypatch.delenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", raising=False) + monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) + result = MiniMaxH3VideoDecodingStage(VAE()).forward( + batch, + SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ), + ) + + assert observed["dtype"] == torch.uint8 + assert torch.equal(result.output, torch.full((1, 3, 5, 16, 16), 63, dtype=torch.uint8)) + + @pytest.mark.parametrize("unified_memory", [True, False]) def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory) -> None: """On unified memory, same-geometry requests reuse one pre-touched CPU buffer; returned frames never alias it. Discrete GPUs allocate per request and hold nothing between them.""" + monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0") latent_shape = (1, 4, 2, 4, 4) rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1)) allocations = [] diff --git a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py index 68e70da3e..917f705cc 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -107,6 +107,22 @@ def test_decode_to_pixels_matches_decode(latent_frames: int) -> None: assert_close(actual, expected, atol=0.0, rtol=0.0) +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs CUDA"))]) +@torch.inference_mode() +def test_decode_to_pixels_uint8_matches_fp32_quantization(device: str) -> None: + torch.manual_seed(20260810) + vae = _tiny_vae().to(device) + latents = torch.randn(1, 4, 12, 4, 4, device=device) + fp32 = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + streamed = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.uint8) + + vae.decode_to_pixels(latents, fp32) + vae.decode_to_pixels(latents, streamed) + + assert torch.equal(streamed, fp32.mul_(255).clamp_(0, 255).to(torch.uint8)) + + # 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))