[perf]: quantize frames to uint8 on-device before the D->H copy

PostDecodeFrameProcessStage was charged ~6s/run (~25% of e2e on Cosmos
2.5, scaling with frames/resolution). Profiling (nsys + microbench)
showed the stage's own compute is only ~0.5-1.9s; the rest is the
non-blocking pinned-CPU samples.copy_(output) D->H of the full fp32
video (~0.3-1.4 GB) completing lazily and blocking the first
postprocess op, plus a single-threaded per-frame CPU *255/cast loop.

Cast to uint8 on the source device (typically CUDA) before the copy:
the transfer becomes 4x smaller (fp32 -> uint8) and the elementwise
work runs on the GPU. Microbench: 2.056s -> 0.034s (T=29), 2.304s ->
0.132s (T=125), 17-60x on the measurable cost.

clamp_(0, 255) additionally fixes a latent overflow: VAE output
slightly outside [0, 1] previously wrapped mod 256 in the unclamped
(x * 255).to(uint8) cast.

Output is not bit-identical to the old CPU cast (float->uint8 differs
<=1 LSB between CPU and GPU on boundary pixels); gate via SSIM rather
than exact equality. Scoped to the pixel-video path only; latent,
audio-only, and return-samples paths are unchanged.

Author: Mister-Raggs

fastvideo/entrypoints/video_generator.py

@@ -792,13 +792,23 @@ def execute_forward_thread():
         if is_latent_output or audio_only:
             frames = None if is_latent_output else []
         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. The previous path read from the pinned-CPU
+            # `samples` buffer whose non-blocking D->H had not completed, so
+            # the first op here blocked on the full fp32 video transfer
+            # (~0.3-1.4 GB) and ran a single-threaded per-frame CPU *255/cast
+            # loop. Casting to uint8 on-device first makes the transfer 4x
+            # smaller 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.
+            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: