From fca45bc8e1b6c7fade69fbb8d9d18987c53c7561 Mon Sep 17 00:00:00 2001 From: Raghav K Date: Thu, 20 Aug 2026 19:55:55 -0700 Subject: [PATCH] [perf] Quantize frames to uint8 on-device before the post-decode D->H copy (#1362) --- fastvideo/entrypoints/video_generator.py | 61 ++++++++---- .../tests/entrypoints/test_video_generator.py | 97 +++++++++++++++++++ 2 files changed, 140 insertions(+), 18 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 789840bc1..f81f338d2 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -808,14 +808,19 @@ class VideoGenerator: latent_batch_size = _infer_latent_batch_size(batch) is_latent_output = fastvideo_args.output_type == "latent" needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output) - needs_samples_buffer = batch.return_frames or needs_frame_output - # When ``output_type == "latent"`` the forward output has latent - # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``) - # rather than the pre-allocation's pixel shape. Skip the pinned - # ~50 MB buffer entirely. Also skip it for metadata-only calls; - # neither the result nor save path will consume the decoded tensor. + # A populated ``samples`` has exactly one consumer — the result + # dict (``"samples": samples if batch.return_frames else None``). + # Post-decode frame building reads ``output_batch.output`` + # directly (the GPU ``vid_u8`` path), not ``samples``. So when + # ``return_frames=False`` the pinned fp32 alloc + D->H copy are + # dead weight — the CLI generate flow (``save_video=True``, + # ``return_frames=False``) hits this on every call. + # ``output_type == "latent"`` keeps its existing branch (shape + # mismatch falls through to ``.cpu()`` below) for callers that + # *do* ask for the latent samples via ``return_frames=True``. # ``skip_pixel_prealloc`` also gates the slow-path warning. - skip_pixel_prealloc = is_latent_output or not needs_samples_buffer + needs_samples_out = batch.return_frames + skip_pixel_prealloc = is_latent_output or not needs_samples_out if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: @@ -835,9 +840,11 @@ class VideoGenerator: "This usually means the executor/pipeline failed earlier.") audio_only = bool(output_batch.extra.get("audio_only")) - if not needs_samples_buffer or (audio_only and not batch.return_frames): - # Metadata-only/audio-only request: keep the empty placeholder and - # avoid the decoded tensor D->H copy. + if not needs_samples_out: + # Nothing downstream reads ``samples`` (the result dict + # returns None when ``return_frames=False``); keep the empty + # placeholder allocated above and skip the fp32 D->H copy + # entirely. pass elif audio_only: # Audio-only return-frames requests expose the small placeholder @@ -869,8 +876,13 @@ class VideoGenerator: # `GenerationResult.size` describes the produced media, not only the # base-stage request. Refiner pipelines can change the final pixel # dimensions, so derive this result metadata from the decoded output. + # Read the geometry from `output_batch.output` (a shape-only access, + # no D->H copy): when `return_frames=False` the `samples` mirror + # stays an empty placeholder and no longer carries the decoded + # shape. Metadata-only calls keep the request fallback and never + # inspect the (possibly dropped) worker output. output_size = _resolve_output_size( - samples, + output_batch.output if needs_frame_output else samples, (target_height, target_width, batch.num_frames), pixel_output=not is_latent_output and not audio_only, ) @@ -882,13 +894,26 @@ class VideoGenerator: elif not needs_frame_output: frames = None else: - videos = rearrange(samples, "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.permute(1, 2, 0).squeeze(-1) - x = (x * 255).to(torch.uint8) - frames.append(x.contiguous().cpu().numpy()) + # Quantize on the source device (typically CUDA) BEFORE the + # device->host copy. `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. The old path + # paid a full fp32 video D->H copy (which scales with + # resolution x frames x batch) and then a single-threaded + # per-frame CPU *255/cast loop. Casting to uint8 on-device + # makes the transfer 4x smaller, ships it in a single copy, + # and moves the elementwise work onto the GPU. clamp_() also + # fixes a latent overflow bug: VAE output slightly outside + # [0, 1] wrapped mod 256 in the old unclamped cast. + # (Equivalence is SSIM-gated, not bit-exact: float->uint8 + # differs <=1 LSB CPU vs GPU.) + src = output_batch.output + vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8) + vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu() + 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/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index d123f60a3..0e7b940fa 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -2,8 +2,10 @@ import os from types import SimpleNamespace import warnings +import numpy as np import pytest import torch +from einops import rearrange import fastvideo.entrypoints.video_generator as video_generator_module from fastvideo.api import ( @@ -271,6 +273,70 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path) assert result["video_path"] is None +def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path): + """The on-device quantize path (#1362) must reproduce the legacy + per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly + for in-range fp32 pixels: same uint8 dtype, same HWC grid layout with + nrow=6 (batch>1), odd frame count. CPU-only: on CUDA the float->uint8 + cast may differ by <=1 LSB, but on CPU both orderings run identical + fp32 ops, so exact equality is required.""" + torch.manual_seed(0) + output = torch.rand((2, 3, 3, 16, 16), dtype=torch.float32) + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + sampling_param = _small_sampling_param(save_video=False, return_frames=True) + sampling_param.num_frames = 3 + sampling_param.num_videos_per_prompt = 2 + + result = generator._generate_single_video( + prompt="grid parity", + sampling_param=sampling_param, + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + legacy_frames = [] + for x in rearrange(output, "b c t h w -> t b c h w"): + grid = video_generator_module.torchvision.utils.make_grid(x, nrow=6) + grid = grid.permute(1, 2, 0).squeeze(-1) + legacy_frames.append((grid * 255).to(torch.uint8).contiguous().cpu().numpy()) + + torch.testing.assert_close(result["samples"], output) + assert len(result["frames"]) == 3 + for got, want in zip(result["frames"], legacy_frames, strict=True): + assert got.dtype == np.uint8 + assert got.shape == want.shape + np.testing.assert_array_equal(got, want) + + +def test_generate_single_video_frames_clamp_out_of_range_pixels(tmp_path): + """VAE output slightly outside [0, 1] must saturate at 0/255 in the + uint8 frames. The pre-#1362 unclamped cast wrapped mod 256 (e.g. + 1.5 -> 126). CPU-only.""" + output = torch.full((1, 3, 2, 16, 16), 1.5, dtype=torch.float32) + output[:, :, 1] = -0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="clamp", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + frames = result["frames"] + assert len(frames) == 2 + # make_grid passes a single image through without grid padding, so + # every pixel comes from the (clamped) output tensor. + assert frames[0].dtype == np.uint8 + assert frames[0].shape == (16, 16, 3) + assert (frames[0] == 255).all() + assert (frames[1] == 0).all() + + def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path): output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 output_batch = _single_video_output_batch(output) @@ -305,6 +371,37 @@ def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_p } +def test_generate_single_video_save_only_reports_refined_output_size(monkeypatch, tmp_path): + """`GenerationResult.size` must describe the decoded media even when the + fp32 `samples` mirror is skipped (`return_frames=False`, the CLI save + flow). Refiner pipelines can change the final pixel geometry, so the size + has to come from `output_batch.output`, not the base request. CPU-only.""" + # Refiner-style output: request asks for 2 frames of 16x16, pipeline + # produces 5 frames of 32x48. + output = torch.full((1, 3, 5, 32, 48), 0.5, dtype=torch.float32) + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + saved = {} + + def fake_mimsave(path, frames, *, fps, format): + saved["frame_count"] = len(frames) + + monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave) + + result = generator._generate_single_video( + prompt="refined save", + sampling_param=_small_sampling_param(save_video=True, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "refined.mp4"), + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["size"] == (32, 48, 5) + assert saved["frame_count"] == 5 + + def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path): audio = torch.zeros((16, ), dtype=torch.float32) output_batch = _single_video_output_batch(