[perf] Quantize frames to uint8 on-device before the post-decode D->H copy (#1362)

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
+        # A populated ``samples`` has exactly one consumer — the result
-        # When ``output_type == "latent"`` the forward output has latent
+        # dict (``"samples": samples if batch.return_frames else None``).
-        # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``)
+        # Post-decode frame building reads ``output_batch.output``
-        # rather than the pre-allocation's pixel shape. Skip the pinned
+        # directly (the GPU ``vid_u8`` path), not ``samples``. So when
-        # ~50 MB buffer entirely. Also skip it for metadata-only calls;
+        # ``return_frames=False`` the pinned fp32 alloc + D->H copy are
-        # neither the result nor save path will consume the decoded tensor.
+        # 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):
+        if not needs_samples_out:
-            # Metadata-only/audio-only request: keep the empty placeholder and
+            # Nothing downstream reads ``samples`` (the result dict
-            # avoid the decoded tensor D->H copy.
+            # 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")
+            # Quantize on the source device (typically CUDA) BEFORE the
-            frames = []
+            # device->host copy. `samples` above is just the pinned-CPU
-            for x in videos:
+            # mirror of `output_batch.output` (`samples.copy_(output)` or
-                x = torchvision.utils.make_grid(x, nrow=6)
+            # `output.cpu()`) with no intervening preprocessing, so reading
-                x = x.permute(1, 2, 0).squeeze(-1)
+            # `output_batch.output` here is the same data. The old path
-                x = (x * 255).to(torch.uint8)
+            # paid a full fp32 video D->H copy (which scales with
-                frames.append(x.contiguous().cpu().numpy())
+            # 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:

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(