[diffusion] improve: improve post-processing by moving compute-intensive tasks to GPU (#15822)

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -12,10 +12,7 @@ import dataclasses
 import os
 
 import imageio
-import numpy as np
 import torch
-import torchvision
-from einops import rearrange
 
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
@@ -66,31 +63,37 @@ def post_process_sample(
     """
     Process sample output and save video if necessary
     """
-    # Process outputs
+    # 1. Vectorized processing on GPU/CPU tensor
     if sample.dim() == 3:
         # for images, dim t is missing
         sample = sample.unsqueeze(1)
-    videos = rearrange(sample, "c t h w -> t c h w")
-    frames = []
-    # TODO: this can be batched
-    for x in videos:
-        x = torchvision.utils.make_grid(x, nrow=6)
-        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-        frames.append((x * 255).numpy().astype(np.uint8))
 
-    # Save outputs if requested
+    # Convert to uint8 and move to CPU in bulk
+    # Shape: [C, T, H, W] -> [T, H, W, C]
+    sample = (sample * 255).clamp(0, 255).to(torch.uint8)
+    videos = sample.permute(1, 2, 3, 0).cpu().numpy()
+
+    # Convert to list of frames for imageio
+    frames = list(videos)
+
+    # 2. Save outputs if requested
     if save_output:
         if save_file_path:
             os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
             if data_type == DataType.VIDEO:
+                # TODO: make this configurable
+                quality = 5
                 imageio.mimsave(
                     save_file_path,
                     frames,
                     fps=fps,
                     format=data_type.get_default_extension(),
+                    codec="libx264",
+                    quality=quality,
                 )
             else:
-                imageio.imwrite(save_file_path, frames[0])
+                quality = 75
+                imageio.imwrite(save_file_path, frames[0], quality=quality)
             logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}")
         else:
             logger.info(f"No output path provided, output not saved")

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -199,23 +199,29 @@ class DecodingStage(PipelineStage):
 
         # decode trajectory latents if needed
         if batch.return_trajectory_decoded:
-            trajectory_decoded = []
             assert (
                 batch.trajectory_latents is not None
             ), "batch should have trajectory latents"
-            for idx in range(batch.trajectory_latents.shape[1]):
-                # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
-                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
-                cur_timestep = batch.trajectory_timesteps[idx]
-                logger.info("decoding trajectory latent for timestep: %s", cur_timestep)
-                decoded_frames = self.decode(cur_latent, server_args)
-                trajectory_decoded.append(decoded_frames.cpu().float())
+
+            # 1. Batch trajectory decoding to improve GPU utilization
+            # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+            B, T, C, F, H, W = batch.trajectory_latents.shape
+            flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W)
+
+            logger.info("decoding %s trajectory latents in batch", B * T)
+            # Use the optimized batch decode
+            all_decoded = self.decode(flat_latents, server_args)
+
+            # 2. Reshape back
+            # Keep on GPU to allow faster vectorized post-processing
+            decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:])
+
+            # Convert to list of tensors (per timestep) as expected by OutputBatch
+            # Each element in list is [B, channels, frames, H_out, W_out]
+            trajectory_decoded = [decoded_tensor[:, i] for i in range(T)]
         else:
             trajectory_decoded = None
 
-        # Convert to CPU float32 for compatibility
-        frames = frames.cpu().float()
-
         # Update batch with decoded image
         output_batch = OutputBatch(
             output=frames,