[CI] [RL]: Add e2e test for custom pipeline (#2005)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

.buildkite/test-ready.yml

@@ -67,6 +67,22 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Custom Pipeline Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Model CPU offloading Test"
     depends_on: upload-ready-pipeline
     commands:

tests/e2e/offline_inference/custom_pipeline/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
+    QwenImagePipelineWithLogProbForTest,
+)
+from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
+    vLLMOmniColocateWorkerExtensionForTest,
+)
+
+__all__ = [
+    "QwenImagePipelineWithLogProbForTest",
+    "vLLMOmniColocateWorkerExtensionForTest",
+]

tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py

@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+@dataclass
+class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            The log probability of the previous sample.
+        prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            The mean of the computed sample of previous timestep.
+        std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
+            The standard deviation used to compute `prev_sample`.
+    """
+
+    prev_sample: torch.FloatTensor
+    log_prob: torch.FloatTensor | None
+    prev_sample_mean: torch.FloatTensor
+    std_dev_t: torch.FloatTensor
+
+
+class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
+    """SDE version of the FlowMatchEulerDiscreteScheduler.
+    The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
+    and diffusers v0.37 branch.
+    """
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+        noise_level: float = 0.7,
+        prev_sample: torch.FloatTensor | None = None,
+        sde_type: Literal["sde", "cps"] = "sde",
+        return_logprobs: bool = True,
+    ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
+            noise_level (`float`, *optional*, defaults to 0.7):
+                The noise level used in the SDE.
+            prev_sample (`torch.FloatTensor`, *optional*):
+                The sample from the previous timestep. If not provided, it will be sampled inside the function.
+            sde_type (`str`, *optional*, defaults to "sde"):
+                The type of SDE to use. Choose between "sde" and "cps".
+            return_logprobs (`bool`, *optional*, defaults to True):
+                Whether to return log probabilities of the previous sample.
+        """
+
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if prev_sample is not None:
+            prev_sample = prev_sample.to(torch.float32)
+
+        prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
+            sample=sample,
+            model_output=model_output,
+            generator=generator,
+            per_token_timesteps=per_token_timesteps,
+            noise_level=noise_level,
+            prev_sample=prev_sample,
+            sde_type=sde_type,
+            return_logprobs=return_logprobs,
+        )
+
+        # upon completion increase step index by one
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if not return_dict:
+            return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
+
+        return FlowMatchSDEDiscreteSchedulerOutputForTest(
+            prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
+        )
+
+    def sample_previous_step(
+        self,
+        sample: torch.Tensor,
+        model_output: torch.Tensor,
+        timestep: torch.FloatTensor | None = None,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        noise_level: float = 0.7,
+        prev_sample: torch.Tensor | None = None,
+        sde_type: Literal["cps", "sde"] = "sde",
+        return_logprobs: bool = True,
+    ):
+        assert sde_type in ["sde", "cps"]
+        assert sample.dtype == torch.float32
+        if prev_sample is not None:
+            assert prev_sample.dtype == torch.float32
+
+        if per_token_timesteps is not None:
+            raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
+        else:
+            if timestep is None:
+                sigma_idx = self.step_index
+                sigma = self.sigmas[sigma_idx]
+                sigma_prev = self.sigmas[sigma_idx + 1]
+            else:
+                sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
+                sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
+                sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
+
+            sigma_max = self.sigmas[1]
+            dt = sigma_prev - sigma
+
+        if sde_type == "sde":
+            std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
+
+            prev_sample_mean = (
+                sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
+                + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
+
+            if return_logprobs:
+                log_prob = (
+                    -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
+                    - torch.log(std_dev_t * torch.sqrt(-1 * dt))
+                    - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
+                )
+            else:
+                log_prob = None
+
+        elif sde_type == "cps":
+            std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
+            pred_original_sample = sample - sigma * model_output
+            noise_estimate = sample + model_output * (1 - sigma)
+            prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
+                sigma_prev**2 - std_dev_t**2
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * variance_noise
+
+            if return_logprobs:
+                log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
+            else:
+                log_prob = None
+
+        # mean along all but batch dimension
+        log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
+        return prev_sample, log_prob, prev_sample_mean, std_dev_t

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -0,0 +1,405 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Detailed custom Qwen-Image pipeline for E2E tests.
+
+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
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Literal
+
+import torch
+
+from tests.e2e.offline_inference.custom_pipeline.flow_match_sde_scheduler import FlowMatchSDEDiscreteSchedulerForTest
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+def _maybe_to_cpu(v):
+    if isinstance(v, torch.Tensor):
+        return v.detach().cpu()
+    return v
+
+
+# Custom pipeline class for QwenImage that returns log probabilities during the diffusion process.
+# This is for test
+class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchSDEDiscreteSchedulerForTest.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder.dtype
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(prompt_ids, dtype=torch.long)
+
+        prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
+        attention_mask = attention_mask.unsqueeze(0) if attention_mask.ndim == 1 else attention_mask
+        drop_idx = self.prompt_template_encode_start_idx
+        encoder_hidden_states = self.text_encoder(
+            input_ids=prompt_ids.to(self.device),
+            attention_mask=attention_mask.to(self.device),
+            output_hidden_states=True,
+        )
+        hidden_states = encoder_hidden_states.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+        return prompt_embeds, encoder_attention_mask
+
+    def encode_prompt(
+        self,
+        prompt_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 1024,
+    ):
+        prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
+        attention_mask = (
+            attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask
+        )
+        batch_size = prompt_ids.shape[0] if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt_ids, attention_mask=attention_mask)
+
+        prompt_embeds = prompt_embeds[:, :max_sequence_length]
+        prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def diffuse(
+        self,
+        prompt_embeds,
+        prompt_embeds_mask,
+        negative_prompt_embeds,
+        negative_prompt_embeds_mask,
+        latents,
+        img_shapes,
+        txt_seq_lens,
+        negative_txt_seq_lens,
+        timesteps,
+        do_true_cfg,
+        guidance,
+        true_cfg_scale,
+        noise_level,
+        sde_window,
+        sde_type,
+        generator,
+        logprobs,
+    ):
+        all_latents = []
+        all_log_probs = []
+        all_timesteps = []
+        self.scheduler.set_begin_index(0)
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            if i < sde_window[0]:
+                cur_noise_level = 0.0
+            elif i == sde_window[0]:
+                cur_noise_level = noise_level
+                all_latents.append(latents)
+            elif i > sde_window[0] and i < sde_window[1]:
+                cur_noise_level = noise_level
+            else:
+                cur_noise_level = 0.0
+
+            self._current_timestep = t
+
+            # Broadcast timestep to match batch size
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            self.transformer.do_true_cfg = do_true_cfg
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+            # Forward pass for negative prompt (CFG)
+            if do_true_cfg:
+                neg_noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    img_shapes=img_shapes,
+                    txt_seq_lens=negative_txt_seq_lens,
+                    attention_kwargs=self.attention_kwargs,
+                    return_dict=False,
+                )[0]
+                comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                noise_pred = comb_pred * (cond_norm / noise_norm)
+            # compute the previous noisy sample x_t -> x_t-1
+            latents, log_prob, _, _ = self.scheduler.step(
+                noise_pred,
+                t,
+                latents,
+                generator=generator,
+                noise_level=cur_noise_level,
+                sde_type=sde_type,
+                return_logprobs=logprobs,
+                return_dict=False,
+            )
+
+            if i >= sde_window[0] and i < sde_window[1]:
+                all_latents.append(latents)
+                all_log_probs.append(log_prob)
+                all_timesteps.append(t)
+
+        all_latents = torch.stack(all_latents, dim=1)
+
+        if all_log_probs[0] is not None:
+            all_log_probs = torch.stack(all_log_probs, dim=1)
+        else:
+            all_log_probs = None
+
+        all_timesteps = torch.stack(all_timesteps).unsqueeze(0).expand(latents.shape[0], -1)
+
+        return latents, all_latents, all_log_probs, all_timesteps
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt_ids: torch.Tensor | list[int] | None = None,
+        prompt_mask: torch.Tensor | None = None,
+        negative_prompt_ids: torch.Tensor | list[int] | None = None,
+        negative_prompt_mask: torch.Tensor | None = None,
+        true_cfg_scale: float = 4.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 1.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds_mask: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end_tensor_inputs: tuple[str, ...] = ("latents",),
+        max_sequence_length: int = 512,
+        noise_level: float = 0.7,
+        sde_window_size: int | None = None,
+        sde_window_range: tuple[int, int] = (0, 5),
+        sde_type: Literal["sde", "cps"] = "sde",
+        logprobs: bool = True,
+    ) -> DiffusionOutput:
+        # Extract prompt data from OmniCustomPrompt in req.prompts[0]
+        custom_prompt = req.prompts[0] if req.prompts else {}
+        if isinstance(custom_prompt, dict):
+            prompt_ids = custom_prompt.get("prompt_ids", prompt_ids)
+            prompt_mask = custom_prompt.get("prompt_mask", prompt_mask)
+            negative_prompt_ids = custom_prompt.get("negative_prompt_ids", negative_prompt_ids)
+            negative_prompt_mask = custom_prompt.get("negative_prompt_mask", negative_prompt_mask)
+
+        # Read sampling params from req.sampling_params
+        sp = req.sampling_params
+        height = sp.height or self.default_sample_size * self.vae_scale_factor
+        width = sp.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = sp.num_inference_steps or num_inference_steps
+        max_sequence_length = sp.max_sequence_length or max_sequence_length
+
+        noise_level = sp.extra_args.get("noise_level", None) or noise_level
+        sde_window_size = sp.extra_args.get("sde_window_size", None) or sde_window_size
+        sde_window_range = sp.extra_args.get("sde_window_range", None) or sde_window_range
+        sde_type = sp.extra_args.get("sde_type", None) or sde_type
+        logprobs = sp.extra_args.get("logprobs", None) or logprobs
+
+        generator = sp.generator or generator
+        if generator is None and sp.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(sp.seed)
+        true_cfg_scale = sp.true_cfg_scale or true_cfg_scale
+        req_num_outputs = getattr(sp, "num_outputs_per_prompt", None)
+        if req_num_outputs and req_num_outputs > 0:
+            num_images_per_prompt = req_num_outputs
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        if prompt_ids is not None:
+            if isinstance(prompt_ids, list):
+                prompt_ids = torch.tensor(prompt_ids, device=self.device)
+            batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+        elif prompt_embeds is not None:
+            batch_size = prompt_embeds.shape[0]
+        else:
+            # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
+            # Return a minimal dummy output to avoid crashing.
+            return DiffusionOutput(output=None, custom_output={})
+
+        if isinstance(negative_prompt_ids, list):
+            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
+
+        has_neg_prompt = negative_prompt_ids is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt_ids=prompt_ids,
+            attention_mask=prompt_mask,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt_ids=negative_prompt_ids,
+                attention_mask=negative_prompt_mask,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        num_channels_latents = self.transformer.in_channels // 4
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+        img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
+
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
+
+        if sde_window_size is not None:
+            start = torch.randint(
+                sde_window_range[0],
+                sde_window_range[1] - sde_window_size + 1,
+                (1,),
+                generator=generator,
+                device=self.device,
+            ).item()
+            end = start + sde_window_size
+            sde_window = (start, end)
+        else:
+            sde_window = (0, len(timesteps) - 1)
+
+        latents, all_latents, all_log_probs, all_timesteps = self.diffuse(
+            prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            latents,
+            img_shapes,
+            txt_seq_lens,
+            negative_txt_seq_lens,
+            timesteps,
+            do_true_cfg,
+            guidance,
+            true_cfg_scale,
+            noise_level,
+            sde_window,
+            sde_type,
+            generator,
+            logprobs,
+        )
+
+        self._current_timestep = None
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+        return DiffusionOutput(
+            output=_maybe_to_cpu(image),
+            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),
+                "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
+            },
+        )

tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Regression tests for AsyncOmni collective_rpc in inline diffusion mode.
+
+When AsyncOmni runs a single diffusion stage it activates "inline diffusion
+mode", which skips stage worker subprocess creation and therefore never
+attaches IPC queues (_in_q / _out_q) to the OmniStage.  Methods like
+list_loras(), add_lora(), sleep(), wake_up() all delegate to
+collective_rpc(), which must handle this mode correctly instead of
+trying to use the non-existent queues.
+
+This is the same code path that verl's vLLMOmniHttpServer.generate()
+exercises when it calls ``await self.engine.list_loras()`` before
+dispatching a generation request.
+
+Usage:
+    pytest tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py -v -s
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import ExitStack
+
+import pytest
+
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+MODEL = "tiny-random/Qwen-Image"
+CUSTOM_PIPELINE_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
+)
+WORKER_EXTENSION_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
+)
+
+
+def _create_inline_engine() -> AsyncOmni:
+    """Create an AsyncOmni instance that uses inline diffusion mode.
+
+    A single diffusion stage triggers inline mode automatically.
+    """
+    engine = AsyncOmni(
+        model=MODEL,
+        custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+        worker_extension_cls=WORKER_EXTENSION_CLASS,
+        enforce_eager=True,
+    )
+
+    return engine
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_list_loras_inline_mode():
+    """list_loras() must not crash in inline diffusion mode.
+
+    This is the exact call that vLLMOmniHttpServer.generate() makes
+    before every generation request.
+    """
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        result = await engine.list_loras()
+        assert isinstance(result, list), f"Expected list, got {type(result)}"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_collective_rpc_inline_mode():
+    """collective_rpc() must delegate to the inline engine, not stage queues."""
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        result = await engine.collective_rpc(method="list_loras")
+        assert isinstance(result, list), f"Expected list, got {type(result)}"
+        assert len(result) == 1, "Inline mode has exactly one stage"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_sleep_wake_up_inline_mode():
+    """sleep() and wake_up() must work in inline diffusion mode."""
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        await engine.sleep(level=1)
+        assert await engine.is_sleeping()
+
+        await engine.wake_up()
+        assert not await engine.is_sleeping()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_generate_after_list_loras_inline_mode():
+    """Full flow: list_loras() then generate(), matching vLLMOmniHttpServer.
+
+    This reproduces the exact sequence that caused the original crash:
+    1. list_loras() (was crashing with AssertionError on _out_q)
+    2. generate() (should succeed)
+    """
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        # Step 1: list_loras (the call that was crashing)
+        loras = await engine.list_loras()
+        assert isinstance(loras, list)
+
+        # Step 2: generate (should still work after list_loras)
+        sampling_params = OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            guidance_scale=0.0,
+            height=256,
+            width=256,
+            seed=42,
+        )
+
+        last_output = None
+        async for output in engine.generate(
+            prompt={"prompt_ids": list(range(50))},
+            request_id=f"test_after_lora_{uuid.uuid4().hex[:8]}",
+            sampling_params_list=[sampling_params],
+            output_modalities=["image"],
+        ):
+            last_output = output
+
+        assert last_output is not None
+        assert isinstance(last_output, OmniRequestOutput)
+        assert last_output.images, "Expected at least one generated image"

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -0,0 +1,239 @@
+# 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)."""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from contextlib import ExitStack
+from functools import lru_cache
+
+import numpy as np
+import pytest
+from transformers import AutoTokenizer
+
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+CUSTOM_PIPELINE_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
+)
+WORKER_EXTENSION_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
+)
+
+# Use your specified HF repo name/path directly here
+MODEL = "tiny-random/Qwen-Image"
+
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+
+
+# ---------------------------------------------------------------------
+#                👇 Test Helper Functions & Fixtures 👇
+# ---------------------------------------------------------------------
+
+_MIN_PROMPT_TOKENS = 35
+
+
+def normalize_token_ids(tokenized_output) -> list[int]:
+    """Normalize tokenizer outputs into a flat ``list[int]``.
+
+    This handles Transformers 4/5 differences where ``apply_chat_template(tokenize=True)``
+    may return either ``list[int]`` or a ``BatchEncoding``/mapping with ``input_ids``.
+    """
+
+    token_ids = tokenized_output
+    if isinstance(tokenized_output, dict):
+        if "input_ids" in tokenized_output:
+            token_ids = tokenized_output["input_ids"]
+    elif hasattr(tokenized_output, "input_ids"):
+        token_ids = tokenized_output.input_ids
+
+    if hasattr(token_ids, "tolist"):
+        token_ids = token_ids.tolist()
+
+    if isinstance(token_ids, tuple):
+        token_ids = list(token_ids)
+
+    if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], list | tuple):
+        token_ids = list(token_ids[0])
+
+    if not isinstance(token_ids, list):
+        raise TypeError(f"token_ids must be list-like token ids, got {type(token_ids).__name__}: {token_ids!r}")
+
+    normalized_ids = []
+    for idx, token_id in enumerate(token_ids):
+        if hasattr(token_id, "item"):
+            token_id = token_id.item()
+        try:
+            normalized_ids.append(int(token_id))
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"token_id must be int-convertible, got {type(token_id).__name__}: {token_id!r}") from e
+    return normalized_ids
+
+
+@lru_cache(maxsize=1)
+def _tokenize_prompt(text: str) -> list[int]:
+    """Tokenize a text prompt into valid token IDs for the model."""
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+    messages = [{"role": "user", "content": text}]
+    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+        f"The pipeline drops the first 34 chat‑template prefix tokens; "
+        f"use a longer prompt so content tokens remain after the drop."
+    )
+    return token_ids
+
+
+def _sampling_params(*, logprobs: bool = False, seed: int = 42) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        height=256,
+        width=256,
+        seed=seed,
+        extra_args={"logprobs": logprobs},
+    )
+
+
+async def _generate_once(
+    engine: AsyncOmni,
+    prompt: str | list[int] | dict,
+    *,
+    request_id: str,
+    sampling_params: OmniDiffusionSamplingParams,
+) -> OmniRequestOutput:
+    # Convert text prompt to dict with tokenized prompt_ids
+    if isinstance(prompt, str):
+        prompt_ids = _tokenize_prompt(prompt)
+        prompt = {"prompt_ids": prompt_ids}
+    elif isinstance(prompt, list):
+        prompt = {"prompt_ids": prompt}
+    # else: assume it's already a dict with prompt_ids
+
+    last_output = None
+    async for output in engine.generate(
+        prompt=prompt,
+        request_id=request_id,
+        sampling_params_list=[sampling_params],
+        output_modalities=["image"],
+    ):
+        last_output = output
+
+    assert last_output is not None
+    assert isinstance(last_output, OmniRequestOutput)
+    return last_output
+
+
+def _assert_valid_image_output(output: OmniRequestOutput) -> None:
+    assert output.final_output_type == "image"
+    assert output.images, "Expected at least one generated image"
+
+    image = output.images[0]
+    arr = np.asarray(image, dtype=np.float32) / 255.0
+
+    assert arr.ndim == 3 and arr.shape[2] == 3, f"Expected HWC RGB image, got shape={arr.shape}"
+    assert arr.shape[0] > 0 and arr.shape[1] > 0
+    assert 0.0 <= float(arr[0, 0, 0]) <= 1.0
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+        )
+        after.callback(engine.shutdown)
+
+        output = await _generate_once(
+            engine,
+            "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+            "reflecting on the calm water surface near a rocky coastline",
+            request_id=f"test_{uuid.uuid4().hex[:8]}",
+            sampling_params=_sampling_params(logprobs=False, seed=42),
+        )
+
+        _assert_valid_image_output(output)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate_with_logprobs():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+        )
+        after.callback(engine.shutdown)
+
+        output = await _generate_once(
+            engine,
+            "a futuristic city at night with neon lights glowing on tall glass "
+            "skyscrapers and flying vehicles soaring between the buildings",
+            request_id=f"test_lp_{uuid.uuid4().hex[:8]}",
+            sampling_params=_sampling_params(logprobs=True, seed=123),
+        )
+
+        _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
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate_concurrent():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+        )
+        after.callback(engine.shutdown)
+
+        prompts = [
+            "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+            "reflecting on the calm water surface near a rocky coastline",
+            "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+            "a garden full of colorful flowers on a bright sunny afternoon",
+            "a majestic mountain landscape covered with fresh white snow under a "
+            "clear blue sky with pine trees in the foreground and a frozen lake",
+            "a futuristic city at night with neon lights glowing on tall glass "
+            "skyscrapers and flying vehicles soaring between the buildings",
+        ]
+
+        tasks = [
+            _generate_once(
+                engine,
+                prompt,
+                request_id=f"concurrent_{i}_{uuid.uuid4().hex[:8]}",
+                sampling_params=_sampling_params(logprobs=False, seed=100 + i),
+            )
+            for i, prompt in enumerate(prompts)
+        ]
+
+        outputs = await asyncio.gather(*tasks)
+
+        assert len(outputs) == len(prompts)
+        for output in outputs:
+            _assert_valid_image_output(output)

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+
+from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
+    vLLMOmniColocateWorkerExtensionForTest,
+)
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_worker_extension_inheritance():
+    assert issubclass(vLLMOmniColocateWorkerExtensionForTest, CustomPipelineWorkerExtension)
+
+
+def test_worker_extension_test_function():
+    assert (
+        vLLMOmniColocateWorkerExtensionForTest.test_extension_name() == "vllm-omni-colocate-worker-extension-for-test"
+    )

tests/e2e/offline_inference/custom_pipeline/worker_extension.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Test-only worker extension for custom diffusion pipeline E2E tests."""
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+
+
+class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
+    """Minimal worker extension used by tests.
+
+    This intentionally stays lightweight: we only inherit the base custom
+    pipeline extension and add one test function.
+    """
+
+    @staticmethod
+    def test_extension_name() -> str:
+        """Return a stable identifier for assertions in unit tests."""
+        return "vllm-omni-colocate-worker-extension-for-test"