[Refactor] Use trajectory_* fields for Qwen-Image structured RL outputs (#2513)

Signed-off-by: samithuang <285365963@qq.com>

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -6,7 +6,8 @@
 This pipeline follows the structure of the user's reference implementation:
 - supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
 - uses an SDE scheduler that can return step logprobs
-- returns rich custom_output fields for testing
+- returns structured trajectory_* fields (latents, timesteps, log_probs)
+  consistent with the BAGEL trajectory recording design
 """
 
 from __future__ import annotations
@@ -393,10 +394,10 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
 
         return DiffusionOutput(
             output=_maybe_to_cpu(image),
+            trajectory_latents=_maybe_to_cpu(all_latents),
+            trajectory_log_probs=_maybe_to_cpu(all_log_probs),
+            trajectory_timesteps=_maybe_to_cpu(all_timesteps),
             custom_output={
-                "all_latents": _maybe_to_cpu(all_latents),
-                "all_log_probs": _maybe_to_cpu(all_log_probs),
-                "all_timesteps": _maybe_to_cpu(all_timesteps),
                 "prompt_embeds": _maybe_to_cpu(prompt_embeds),
                 "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
                 "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -1,7 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""E2E tests for AsyncOmni Qwen-Image generation flow (no Ray, no HTTP server)."""
+"""E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields.
+
+Validates that the custom Qwen-Image pipeline returns structured trajectory
+outputs (latents, timesteps, log_probs) via OmniRequestOutput's trajectory_*
+fields instead of the legacy custom_output dict.
+"""
 
 from __future__ import annotations
 
@@ -191,10 +196,17 @@ async def test_async_omni_generate_with_logprobs():
 
         _assert_valid_image_output(output)
 
-        all_log_probs = output.custom_output.get("all_log_probs")
-        assert all_log_probs is not None, "all_log_probs should be present when logprobs=True"
-        assert hasattr(all_log_probs, "shape")
-        assert all_log_probs.numel() > 0
+        assert output.trajectory_latents is not None, "trajectory_latents should be present"
+        assert hasattr(output.trajectory_latents, "shape")
+        assert output.trajectory_latents.numel() > 0
+
+        assert output.trajectory_timesteps is not None, "trajectory_timesteps should be present"
+        assert hasattr(output.trajectory_timesteps, "shape")
+        assert output.trajectory_timesteps.numel() > 0
+
+        assert output.trajectory_log_probs is not None, "trajectory_log_probs should be present when logprobs=True"
+        assert hasattr(output.trajectory_log_probs, "shape")
+        assert output.trajectory_log_probs.numel() > 0
 
 
 @pytest.mark.core_model

vllm_omni/diffusion/diffusion_engine.py

@@ -196,6 +196,10 @@ class DiffusionEngine:
                         prompt=prompt,
                         metrics=metrics,
                         latents=output.trajectory_latents,
+                        trajectory_latents=output.trajectory_latents,
+                        trajectory_timesteps=output.trajectory_timesteps,
+                        trajectory_log_probs=output.trajectory_log_probs,
+                        trajectory_decoded=output.trajectory_decoded,
                         multimodal_output={"audio": request_audio_payload},
                         final_output_type="audio",
                         stage_durations=output.stage_durations,
@@ -252,6 +256,10 @@ class DiffusionEngine:
                             prompt=prompt,
                             metrics=metrics,
                             latents=output.trajectory_latents,
+                            trajectory_latents=output.trajectory_latents,
+                            trajectory_timesteps=output.trajectory_timesteps,
+                            trajectory_log_probs=output.trajectory_log_probs,
+                            trajectory_decoded=output.trajectory_decoded,
                             multimodal_output={"audio": request_audio_payload},
                             final_output_type="audio",
                             stage_durations=output.stage_durations,

vllm_omni/outputs.py

@@ -123,8 +123,12 @@ class OmniRequestOutput:
             prompt: The prompt used
             metrics: Generation metrics
             latents: Optional latent tensors
+            trajectory_latents: Optional stacked trajectory latent tensors
+            trajectory_timesteps: Optional stacked trajectory timestep tensors
+            trajectory_log_probs: Optional stacked trajectory log-probability tensors
+            trajectory_decoded: Optional list of decoded trajectory images
             multimodal_output: Optional multimodal output dict
-            custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds)
+            custom_output: Optional custom output dict (e.g. prompt embeds)
             stage_durations: Optional stage durations (execution time of each stage) dict
             peak_memory_mb: Peak memory usage in MB