[Feature] Add trajectory recording to BAGEL denoising loop (#2483)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Signed-off-by: Zhengyuan Su (苏政渊) <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Samit <285365963@qq.com>

vllm_omni/diffusion/ipc.py

@@ -78,13 +78,29 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     return tensor
 
 
+def _pack_tensor_if_large(val: torch.Tensor) -> torch.Tensor | dict:
+    """Replace a tensor with an SHM handle if it exceeds the threshold."""
+    if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD:
+        return _tensor_to_shm(val)
+    return val
+
+
+def _unpack_if_shm_handle(val: object) -> object:
+    """Reconstruct a tensor from an SHM handle dict, or return as-is."""
+    if isinstance(val, dict) and val.get("__tensor_shm__"):
+        return _tensor_from_shm(val)
+    return val
+
+
 def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     if output.output is not None and isinstance(output.output, torch.Tensor):
-        if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
-            output.output = _tensor_to_shm(output.output)
+        output.output = _pack_tensor_if_large(output.output)
     if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
-        if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
-            output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
+        output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents)
+    if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor):
+        output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps)
+    if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
+        output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs)
     return output
 
 
@@ -104,10 +120,10 @@ def pack_diffusion_output_shm(output: object) -> object:
 
 
 def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
-    if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
-        output.output = _tensor_from_shm(output.output)
-    if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
-        output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
+    output.output = _unpack_if_shm_handle(output.output)
+    output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents)
+    output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps)
+    output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs)
     return output