nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Raghav K
2026-08-20 21:55:55 -05:00
committed by GitHub
parent 86d639c848
commit fca45bc8e1
2 changed files with 140 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/entrypoints/video_generator.py
+43 -18
View File
@@ -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:
fastvideo/tests/entrypoints/test_video_generator.py
+97
View File
@@ -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(