diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 3719b5d3c..303be8c8a 100644 --- a/.buildkite/test-ready.yml +++ b/.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: diff --git a/tests/e2e/offline_inference/custom_pipeline/__init__.py b/tests/e2e/offline_inference/custom_pipeline/__init__.py new file mode 100644 index 000000000..99663a46a --- /dev/null +++ b/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", +] diff --git a/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py b/tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py new file mode 100644 index 000000000..f85eb6602 --- /dev/null +++ b/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 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 new file mode 100644 index 000000000..ed5b219f8 --- /dev/null +++ b/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), + }, + ) diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py new file mode 100644 index 000000000..57743d62b --- /dev/null +++ b/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" 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 new file mode 100644 index 000000000..f1b4595c9 --- /dev/null +++ b/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) diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py new file mode 100644 index 000000000..b809bed10 --- /dev/null +++ b/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" + ) diff --git a/tests/e2e/offline_inference/custom_pipeline/worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/worker_extension.py new file mode 100644 index 000000000..962bc0127 --- /dev/null +++ b/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"