From e73245c60fca55eba5e86b2704e8e266e88ad93b Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Sat, 4 Jul 2026 11:45:49 +0800 Subject: [PATCH] [Bugfix] Adapt ZImagePipeline.forward calling in MingImagePipeline with DiffusionRequestBatch (#4866) Signed-off-by: Yuanheng Zhao --- .../ming_flash_omni/pipeline_ming_imagegen.py | 89 ++++++++----------- 1 file changed, 39 insertions(+), 50 deletions(-) diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py index f3e0f3357..c67f8cb24 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py @@ -24,7 +24,6 @@ from __future__ import annotations import logging import os -from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -49,7 +48,9 @@ from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import ( MingZImageTransformer2DModel, ) from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -58,38 +59,6 @@ from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenCon logger = logging.getLogger(__name__) -@dataclass -class _ZPipelineSamplingParams: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req.sampling_params``. Unlike SimpleNamespace this fails loudly - at construction if a required field is missing.""" - - height: int - width: int - num_inference_steps: int - guidance_scale: float - generator: torch.Generator | None = None - strength: float | None = None - sigmas: list[float] | None = None - max_sequence_length: int = 512 - guidance_rescale: float | None = None - num_outputs_per_prompt: int = 1 - - -@dataclass -class _ZPipelineRequest: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req``.""" - - request_id: str - sampling_params: _ZPipelineSamplingParams - prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}]) - - @property - def num_reqs(self) -> int: - return 1 - - class MingImagePipeline(ZImagePipeline): """Ming-flash-omni-2.0 text-to-image diffusion pipeline. @@ -123,6 +92,10 @@ class MingImagePipeline(ZImagePipeline): self.device = self._execution_device # Ming convention alias self._dtype = dtype + # Request-scoped conditioning handed to the inherited encode_prompt override + self._pending_prompt_embeds: list[torch.Tensor] | None = None + self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None + # Ming's per-checkpoint image-gen configuration. We cannot rely on # ``od_config.hf_config.image_gen_config`` because the diffusion # stage is started with ``hf_config_name: thinker_config`` (the @@ -273,6 +246,21 @@ class MingImagePipeline(ZImagePipeline): latent = self.vae.encode(ref).latent_dist.mode() return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor + def encode_prompt(self, *args, **kwargs): # noqa: ARG002 + """Return Ming's precomputed conditioning instead of encoding text. + + NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented) + is computed in forward and stashed on `self._pending_*` immediately before + the `super().forward` call, so we simply hand it back here. + """ + if self._pending_prompt_embeds is None: + raise RuntimeError( + "MingImagePipeline.encode_prompt called without pending " + "conditioning; it must run within MingImagePipeline.forward." + ) + + return self._pending_prompt_embeds, self._pending_negative_prompt_embeds + # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ @@ -416,17 +404,27 @@ class MingImagePipeline(ZImagePipeline): negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])] else: negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds] + self._pending_prompt_embeds = prompt_embeds + self._pending_negative_prompt_embeds = negative_prompt_embeds - z_sp = _ZPipelineSamplingParams( + # Build a real single-request DiffusionRequestBatch; + # the prompt text itself is a neutral placeholder here. + z_sp = OmniDiffusionSamplingParams( height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, + output_type="pt", ) - z_req = _ZPipelineRequest( - request_id=req.request_id or "ming-imagegen", - sampling_params=z_sp, + z_req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt={"prompt": ""}, + sampling_params=z_sp, + request_id=req.request_id or "ming-imagegen", + ) + ] ) # Reference image (img2img) → VAE-encoded latent published on the @@ -446,21 +444,12 @@ class MingImagePipeline(ZImagePipeline): None if ref_latent is None else tuple(ref_latent.shape), ) try: - outputs: DiffusionOutput = super().forward( - z_req, - prompt=None, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generator=generator, - output_type="pt", - return_dict=True, - ) + outputs: DiffusionOutput = super().forward(z_req) finally: set_forward_context_ref_latent(None) + # Drop request-scoped conditioning so we don't retain GPU tensors. + self._pending_prompt_embeds = None + self._pending_negative_prompt_embeds = None raw = outputs.output if not isinstance(raw, torch.Tensor):