[Bugfix] Align Ming-Image prompt templates and RNG, fix number of layer to decompose (#8118)

Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>

recipes/inclusionAI/Ming-Image.md

@@ -27,7 +27,7 @@ For Design-Layer, the first returned image is the reconstructed composite and th
 - Python: 3.10+
 - CUDA: 13.0
 - vLLM version: 0.29.0
-- vLLM-Omni version or commit: 59db6a428
+- vLLM-Omni version or commit: 6daf5b30f
 - 2x H100 80GB (1xH100 to be validated)
 
 ### Commands
@@ -52,7 +52,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
     "modalities": ["image"],
     "extra_body": {
       "height": 1024, "width": 1024,
-      "steps": 12, "cfg": 1.0, "seed": 42
+      "num_inference_steps": 12, "guidance_scale": 1.0, "seed": 42
     }
   }' \
   | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
@@ -79,7 +79,7 @@ jq -n \
     modalities: ["image"],
     extra_body: {
       height: 1024, width: 1024,
-      steps: 12, cfg: 1.0, seed: 42
+      num_inference_steps: 12, guidance_scale: 1.0, seed: 42
     }
   }' |
 curl -sS http://127.0.0.1:8091/v1/chat/completions \
@@ -97,7 +97,7 @@ Note that the prompt should better depict each layer to be decomposed, we will r
 ```bash
 MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
 INPUT_IMAGE=/path/to/input.png
-PROMPT="Decompose this design into editable visual layers"
+PROMPT="Decompose this image into 6 layers with the following specifications: Number of layers: 6 \nLayer 1: Central text ... Layer 2: ... Layer 6: ..."
 
 jq -n \
   --arg model "$MODEL" \
@@ -113,7 +113,7 @@ jq -n \
     extra_body: {
       num_layers: 6,
       height: 1024, width: 1024,
-      steps: 12, cfg: 2.0, seed: 42
+      num_inference_steps: 12, guidance_scale: 2.0, seed: 42
     }
   }' |
 curl -sS http://127.0.0.1:8091/v1/chat/completions \

tests/model_executor/models/ming_image/test_config_and_pipeline.py

@@ -5,20 +5,27 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.multimodal.processing import ProcessorInputs
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.ming_image.pipeline import (
     MingImageDiffusionPipeline,
     _validate_variant_config,
 )
+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.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.models.ming_flash_omni.ming_flash_omni_thinker import (
     MingFlashOmniThinkerForConditionalGeneration,
+    MingFlashOmniThinkerMultiModalProcessor,
 )
 from vllm_omni.model_executor.models.ming_image import checkpoint
-from vllm_omni.model_executor.models.ming_image.model import MingImageForConditionalGeneration
+from vllm_omni.model_executor.models.ming_image.model import (
+    MingImageForConditionalGeneration,
+    MingImageMultiModalProcessor,
+)
 from vllm_omni.model_executor.models.ming_image.pipeline import MING_IMAGE_PIPELINE
 from vllm_omni.model_executor.stage_input_processors.ming_image import thinker2image
 from vllm_omni.model_extras import (
@@ -26,6 +33,7 @@ from vllm_omni.model_extras import (
     should_init_extra_args_for_non_diffusion_stages,
 )
 from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
+from vllm_omni.transformers_utils.processors.ming import MingImageProcessor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -119,6 +127,21 @@ def test_pipeline_rejects_multiple_outputs_per_prompt():
         pipeline.forward(request)
 
 
+@pytest.mark.parametrize(
+    ("guidance_scale", "ming_expected", "z_image_expected"),
+    ((0.0, False, False), (1.0, False, True), (1.0001, True, True), (2.0, True, True)),
+)
+def test_ming_cfg_is_only_enabled_above_one(guidance_scale, ming_expected, z_image_expected):
+    # Ming-Image follows the vendor threshold; the shared Z-Image base keeps diffusers semantics.
+    ming_pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    ming_pipeline._guidance_scale = guidance_scale
+    z_image_pipeline = ZImagePipeline.__new__(ZImagePipeline)
+    z_image_pipeline._guidance_scale = guidance_scale
+
+    assert ming_pipeline.do_classifier_free_guidance is ming_expected
+    assert z_image_pipeline.do_classifier_free_guidance is z_image_expected
+
+
 def _source_output(prefix_len: int = 4):
     prompt_ids = [11] * prefix_len + [157158] + [157157] * 256 + [157159]
     length = len(prompt_ids)
@@ -157,6 +180,78 @@ def test_bridge_rejects_text_negative_condition():
         )
 
 
+def test_image_generation_template_matches_upstream_layout():
+    processor = MingImageProcessor.__new__(MingImageProcessor)
+    processor.tokenizer = SimpleNamespace(eos_token="<eos>")
+
+    text = processor._apply_image_generation_template("draw a cabin")
+    assert text == (
+        "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off"
+        "<eos><role>HUMAN</role>draw a cabin"
+        "<eos><role>ASSISTANT</role>"
+    )
+    assert processor._apply_image_generation_template(text) == text
+    assert "<role>HUMAN</role><IMAGE>edit this image<eos>" in processor._apply_image_generation_template(
+        "edit this image",
+        has_reference_image=True,
+    )
+
+
+def test_multimodal_processor_applies_generation_template(monkeypatch: pytest.MonkeyPatch):
+    template_calls = []
+
+    def _format(prompt, *, has_reference_image):
+        template_calls.append((prompt, has_reference_image))
+        return "formatted prompt"
+
+    tokenizer = SimpleNamespace(
+        decode=lambda token_ids, **kwargs: "edit this image",
+        encode=lambda text, **kwargs: [101, 102],
+    )
+    processor = MingImageMultiModalProcessor.__new__(MingImageMultiModalProcessor)
+    processor.info = SimpleNamespace(
+        get_tokenizer=lambda: tokenizer,
+        get_hf_processor=lambda: SimpleNamespace(_apply_image_generation_template=_format),
+    )
+    captured = {}
+
+    def _parent_apply(_self, inputs, timing_ctx):
+        captured["inputs"] = inputs
+        return "processed"
+
+    monkeypatch.setattr(MingFlashOmniThinkerMultiModalProcessor, "apply", _parent_apply)
+    inputs = ProcessorInputs(
+        prompt=[11, 12],
+        mm_data_items=None,
+        hf_processor_mm_kwargs={"modalities": ["img2img"]},
+    )
+
+    assert processor.apply(inputs, None) == "processed"
+    assert template_calls == [("edit this image", True)]
+    assert captured["inputs"].prompt == [101, 102]
+    assert captured["inputs"].hf_processor_mm_kwargs["modalities"] == ["image"]
+
+
+@pytest.mark.parametrize(
+    ("pipeline_config", "expected_class_name"),
+    ((MING_IMAGE_PIPELINE, "MingImageDiffusionPipeline"), (None, None)),
+)
+def test_engine_od_config_falls_back_to_registered_pipeline_class(monkeypatch, pipeline_config, expected_class_name):
+    # Ming-Image ships neither a root config.json nor model_index.json for now,
+    # so class-name resolution returns None; serving must still find the extras (e.g. num_layers).
+    monkeypatch.setattr("vllm_omni.diffusion.data.resolve_model_class_name", lambda model, **kwargs: None)
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "inclusionAI/Ming-Image-0.1-Design-Layer"
+    engine.pipeline_config = pipeline_config
+    engine._diffusion_od_config_view = None
+
+    od_config = engine.get_diffusion_od_config()
+
+    assert od_config.model_class_name == expected_class_name
+    assert ("num_layers" in get_extra_body_params(od_config.model_class_name)) is (expected_class_name is not None)
+    assert engine.get_diffusion_od_config() is od_config
+
+
 def test_two_stage_topology_and_request_metadata():
     assert MING_IMAGE_PIPELINE.model_type == "ming_image"
     assert [stage.model_stage for stage in MING_IMAGE_PIPELINE.stages] == ["mllm", "dit"]
@@ -168,6 +263,8 @@ def test_two_stage_topology_and_request_metadata():
     assert MING_IMAGE_PIPELINE.hf_architectures == ("MingImageForConditionalGeneration",)
     for class_name in MING_IMAGE_PIPELINE.diffusers_class_aliases + (MING_IMAGE_PIPELINE.diffusers_class_name,):
         assert "num_layers" in get_extra_body_params(class_name)
+        assert "steps" not in get_extra_body_params(class_name)
+        assert "cfg" not in get_extra_body_params(class_name)
         assert should_init_extra_args_for_non_diffusion_stages(class_name)
 
 

vllm_omni/config/stage_config.py

@@ -312,10 +312,10 @@ class PipelineConfig:
     # ``hf_config_predicate=lambda c: getattr(c, "version", "") == "4.5"``
     # to avoid misrouting 2.6 checkpoints.
     hf_config_predicate: Callable[[Any], bool] | None = None
-    # Diffusers pipeline class name: for models that ship a ``model_index.json``
-    # (no root ``config.json``), the ``_class_name`` field is matched against
-    # this value to auto-detect the pipeline.  Only needed for diffusers-style
-    # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
+    # Canonical diffusion runtime class for this registered pipeline.
+    # Serving uses it as a metadata fallback; model_index.json discovery
+    # also matches its _class_name against this value and its aliases.
+    # None delegates class discovery to the checkpoint metadata.
     diffusers_class_name: str | None = None
     diffusers_class_aliases: tuple[str, ...] = ()
     endpoint_restrictions: tuple[EndpointRestriction, ...] = ()

vllm_omni/diffusion/models/ming_image/pipeline.py

@@ -107,6 +107,11 @@ class MingImageDiffusionPipeline(ZImagePipeline):
     _encoder_modules: ClassVar[list[str]] = ["conditioning"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
 
+    @property
+    def do_classifier_free_guidance(self) -> bool:
+        # The vendor Ming-Image pipeline only enables CFG above 1.0
+        return self._guidance_scale > 1
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         del prefix
         nn.Module.__init__(self)
@@ -190,6 +195,7 @@ class MingImageDiffusionPipeline(ZImagePipeline):
         self.image_processor = VaeImageProcessor(
             vae_scale_factor=self.vae_scale_factor * 2,
             do_convert_rgb=False,
+            resample="bilinear",
         )
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler
@@ -314,18 +320,10 @@ class MingImageDiffusionPipeline(ZImagePipeline):
 
         height = int(extra_args.get("height") or sampling.height or 1024)
         width = int(extra_args.get("width") or sampling.width or 1024)
-        steps = int(extra_args.get("steps") or sampling.num_inference_steps or self.default_num_inference_steps)
-        cfg = float(
-            extra_args["cfg"]
-            if extra_args.get("cfg") is not None
-            else sampling.guidance_scale
-            if sampling.guidance_scale is not None
-            else self.default_guidance_scale
-        )
+        steps = int(sampling.num_inference_steps or self.default_num_inference_steps)
+        cfg = float(sampling.guidance_scale if sampling.guidance_scale is not None else self.default_guidance_scale)
         seed = extra_args.get("seed", sampling.seed)
-        generator = (
-            torch.Generator(device=self.device).manual_seed(int(seed)) if seed is not None else sampling.generator
-        )
+        generator = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else sampling.generator
 
         ref_latent = self._encode_reference(reference, height, width)
         positive = [item for item in cap_feats]
@@ -333,7 +331,7 @@ class MingImageDiffusionPipeline(ZImagePipeline):
         self._pending_prompt_embeds = positive
         self._pending_negative_prompt_embeds = negative
 
-        apply_cfg = cfg > 0
+        apply_cfg = cfg > 1
         context_direct = (
             torch.cat([direct_condition, torch.zeros_like(direct_condition)], dim=0) if apply_cfg else direct_condition
         )

vllm_omni/engine/omni_engine_base.py

@@ -356,6 +356,9 @@ class OmniEngineBase:
             from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
             model_class_name = resolve_model_class_name(self.model)
+            if model_class_name is None:
+                # use registered diffusers cls name as fallback
+                model_class_name = getattr(self.pipeline_config, "diffusers_class_name", None)
             metadata = get_diffusion_model_metadata(model_class_name)
             self._diffusion_od_config_view = SimpleNamespace(
                 model_class_name=model_class_name,

vllm_omni/model_executor/models/ming_image/model.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Mapping
+from dataclasses import replace
 from pathlib import Path
 from typing import Any
 
@@ -88,7 +89,40 @@ class MingImageDummyInputsBuilder(MingFlashOmniThinkerDummyInputsBuilder):
 
 
 class MingImageMultiModalProcessor(MingFlashOmniThinkerMultiModalProcessor):
-    pass
+    def apply(self, inputs, timing_ctx):
+        modalities = list(inputs.hf_processor_mm_kwargs.get("modalities") or [])
+        is_image_generation = "image" in modalities or "img2img" in modalities
+        if is_image_generation:
+            tokenizer = self.info.get_tokenizer()
+            if isinstance(inputs.prompt, str):
+                prompt_text = inputs.prompt
+            else:
+                prompt_text = tokenizer.decode(
+                    inputs.prompt,
+                    skip_special_tokens=False,
+                )
+
+            processor = self.info.get_hf_processor()
+            # expects to place <IMAGE> inside the HUMAN message.
+            formatted_prompt = processor._apply_image_generation_template(
+                prompt_text,
+                has_reference_image="img2img" in modalities,
+            )
+            prompt_ids = tokenizer.encode(
+                formatted_prompt,
+                add_special_tokens=False,
+            )
+
+            # Normalize img2img so the shared parent processor does't prepend a second placeholder outside that message
+            processor_kwargs = dict(inputs.hf_processor_mm_kwargs)
+            processor_kwargs["modalities"] = ["image" if modality == "img2img" else modality for modality in modalities]
+            inputs = replace(
+                inputs,
+                prompt=prompt_ids,
+                hf_processor_mm_kwargs=processor_kwargs,
+            )
+
+        return super().apply(inputs, timing_ctx)
 
 
 @MULTIMODAL_REGISTRY.register_processor(

vllm_omni/model_extras/ming_image.py

@@ -10,8 +10,6 @@ MING_IMAGE_EXTRA_BODY_PARAMS = frozenset(
     {
         "height",
         "width",
-        "steps",
-        "cfg",
         "seed",
         "num_layers",
         "negative_prompt",

vllm_omni/transformers_utils/processors/ming.py

@@ -552,6 +552,19 @@ class MingImageProcessor(ProcessorMixin):
             expanded.append(sample)
         return expanded
 
+    def _apply_image_generation_template(
+        self,
+        prompt: str,
+        *,
+        has_reference_image: bool = False,
+    ) -> str:
+        if prompt.startswith("<role>SYSTEM</role>"):
+            return prompt
+        # Ming-Image uses a fixed single-turn template, so build it directly here
+        image_prefix = PLACEHOLDER_IMAGE_TOKEN_IN_TEXT if has_reference_image else ""
+        eos = self.tokenizer.eos_token
+        return SYSTEM_PROMPT_NOTHINK + eos + USER_PREFIX + image_prefix + prompt + eos + ASSISTANT_PREFIX
+
     def batch_decode(self, *args, **kwargs):
         return self.tokenizer.batch_decode(*args, **kwargs)