From 956f53b2781dde13480b9082fd62fd1a42df1fc2 Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Thu, 9 Apr 2026 18:06:40 +0800 Subject: [PATCH] [Refactor] Use trajectory_* fields for Qwen-Image structured RL outputs (#2513) Signed-off-by: samithuang <285365963@qq.com> --- .../qwen_image_pipeline_with_logprob.py | 9 ++++---- .../test_async_omni_qwen_image_generate.py | 22 ++++++++++++++----- vllm_omni/diffusion/diffusion_engine.py | 8 +++++++ vllm_omni/outputs.py | 6 ++++- 4 files changed, 35 insertions(+), 10 deletions(-) diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py index ed5b219f8..709c66555 100644 --- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py +++ b/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), diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py index f1b4595c9..03bd12efa 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/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 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 5b77c064f..422ef479b 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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, diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 4a775356e..9a7bb6706 100644 --- a/vllm_omni/outputs.py +++ b/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