[New Model] Ming-Image (inclusionAI/Ming-Image-0.1-Design) Support (#8021)

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

recipes/inclusionAI/Ming-Image.md

@@ -0,0 +1,141 @@
+# Ming-Image 0.1 Design
+
+> Text-to-image, image editing, and layer decomposition.
+
+## Summary
+
+- Vendor: inclusionAI
+- Models: `inclusionAI/Ming-Image-0.1-Design` and `inclusionAI/Ming-Image-0.1-Design-Layer`
+- Runtime: vLLM-Omni two-stage online serving
+- API: OpenAI-compatible chat completions
+
+## Architecture
+
+Stage 0 uses a Qwen2.5-VL vision tower plus a 20-layer BailingMoeV2 language model.
+It appends 256 learned image-query tokens and exports the final query states plus direct VLM states from layers 5, 12, and 20.
+The query-token checkpoint is in the root sibling `mlp/` directory, so remote stage-0 materialization downloads both `mllm/` and `mlp/`.
+
+Stage 1 projects those conditions through the checkpoint connector and runs a 30-layer Z-Image DiT with a Qwen-Image RGBA VAE.
+
+For Design-Layer, the first returned image is the reconstructed composite and the remaining images are the requested layers.
+
+## CUDA
+
+### Environment
+
+- OS: Linux
+- Python: 3.10+
+- CUDA: 13.0
+- vLLM version: 0.29.0
+- vLLM-Omni version or commit: 59db6a428
+- 2x H100 80GB (1xH100 to be validated)
+
+### Commands
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design
+vllm serve "$MODEL" --omni --deploy-config vllm_omni/deploy/ming_image.yaml --port 8091
+```
+
+For layer decomposition, set `MODEL` to `inclusionAI/Ming-Image-0.1-Design-Layer`.
+
+## Text-to-image
+
+Note that a prompt refiner is expected to describe the prompts with details; we will refine with more example inputs soon.
+
+```bash
+curl -s http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "inclusionAI/Ming-Image-0.1-Design",
+    "messages": [{"role": "user", "content": "A clean editorial botanical poster"}],
+    "modalities": ["image"],
+    "extra_body": {
+      "height": 1024, "width": 1024,
+      "steps": 12, "cfg": 1.0, "seed": 42
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_design_smoke.png
+```
+
+## Image editing
+
+Pass one input image and an editing instruction:
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design
+INPUT_IMAGE=/path/to/input.png
+
+jq -n \
+  --arg model "$MODEL" \
+  --rawfile image <(base64 -w0 "$INPUT_IMAGE") \
+  '{
+    model: $model,
+    messages: [{role: "user", content: [
+      {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
+      {type: "text", text: "Change the background to blue"}
+    ]}],
+    modalities: ["image"],
+    extra_body: {
+      height: 1024, width: 1024,
+      steps: 12, cfg: 1.0, seed: 42
+    }
+  }' |
+curl -sS http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  --data-binary @- \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_edit.png
+```
+
+## Layer decomposition
+
+Use the Design-Layer checkpoint and set `INPUT_IMAGE` to a local flattened design image.
+Note that the prompt should better depict each layer to be decomposed, we will refine with more example inputs soon.
+
+```bash
+MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
+INPUT_IMAGE=/path/to/input.png
+PROMPT="Decompose this design into editable visual layers"
+
+jq -n \
+  --arg model "$MODEL" \
+  --rawfile image <(base64 -w0 "$INPUT_IMAGE") \
+  --arg prompt "$PROMPT" \
+  '{
+    model: $model,
+    messages: [{role: "user", content: [
+      {type: "image_url", image_url: {url: ("data:image/png;base64," + $image)}},
+      {type: "text", text: $prompt}
+    ]}],
+    modalities: ["image"],
+    extra_body: {
+      num_layers: 6,
+      height: 1024, width: 1024,
+      steps: 12, cfg: 2.0, seed: 42
+    }
+  }' |
+curl -sS http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  --data-binary @- > response.json
+
+jq -r '.choices[0].message.content[].image_url.url | split(",")[1]' response.json |
+  nl -v 0 |
+  while read -r index data; do
+    printf "%s" "$data" | base64 -d > "ming_layer_${index}.png"
+  done
+```
+
+## Notes
+
+- The default deployment keeps Stage 0 eager and enables dynamic regional
+  compilation with CUDA Graph Trees for the repeated Stage 1 DiT blocks.
+  The first request for each new image shape or layer count pays compilation
+  and graph-capture cost; warm up every production shape before measuring or
+  serving latency-sensitive traffic.
+- Only one reference image and one request at a time are currently supported.
+- Design-Layer requires a reference image except during warmup.
+- A non-empty `negative_prompt` is rejected; Ming-Image uses zero negative conditioning.
+- Height and width must be divisible by 16.
+- Returned images retain the checkpoint's four RGBA channels.

tests/diffusion/cache/test_teacache_extractors.py

@@ -33,6 +33,7 @@ from vllm_omni.diffusion.cache.teacache.extractors import (
     extract_flux2_klein_context,
     extract_flux_context,
     extract_minimax_h3_context,
+    extract_zimage_context,
 )
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
@@ -53,6 +54,62 @@ def setup_tp_group():
             yield
 
 
+@pytest.mark.cpu
+def test_zimage_extractor_accepts_extended_patchify_output():
+    class _Block(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attention_norm1 = nn.Identity()
+
+        def adaLN_modulation(self, value):
+            return value.new_zeros((value.shape[0], 16))
+
+    class _Model(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layers = nn.ModuleList([_Block()])
+            self.noise_refiner = nn.ModuleList()
+            self.context_refiner = nn.ModuleList()
+            self.all_x_embedder = nn.ModuleDict({"2-1": nn.Linear(2, 4)})
+            self.cap_embedder = nn.Linear(3, 4)
+            self.x_pad_token = nn.Parameter(torch.zeros(1, 4))
+            self.cap_pad_token = nn.Parameter(torch.zeros(1, 4))
+            self.t_scale = 1.0
+
+        @staticmethod
+        def t_embedder(timestep):
+            return timestep.new_ones((timestep.shape[0], 4))
+
+        @staticmethod
+        def rope_embedder(position_ids):
+            rope = torch.zeros((position_ids.shape[0], 2))
+            return rope, rope
+
+        @staticmethod
+        def patchify_and_embed(*_args):
+            positions = [torch.zeros((32, 3), dtype=torch.int32)]
+            masks = [torch.zeros(32, dtype=torch.bool)]
+            return (
+                [torch.zeros((32, 2))],
+                [torch.zeros((32, 3))],
+                [(1, 1, 1)],
+                positions,
+                positions,
+                masks,
+                masks,
+                [],
+            )
+
+    context = extract_zimage_context(
+        _Model(),
+        x=[torch.zeros((1, 1, 1, 1))],
+        t=torch.zeros(1),
+        cap_feats=[torch.zeros((1, 3))],
+    )
+
+    assert context.modulated_input.shape == (1, 64, 4)
+
+
 class BaseExtractorTest(ABC):
     """Base class for testing TeaCache extractors.
 

tests/diffusion/models/z_image/test_zimage_tp_constraints.py

@@ -1,10 +1,47 @@
-import pytest
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.z_image.z_image_transformer import (
+    UnifiedPrepare,
+    validate_zimage_tp_constraints,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_unified_prepare_applies_padding_masks_before_sharding():
+    x = torch.zeros((2, 4, 2))
+    cap = torch.zeros((2, 3, 2))
+    x_rope = torch.zeros((2, 4, 1))
+    cap_rope = torch.zeros((2, 3, 1))
+    x_mask = torch.tensor([[1, 0, 1, 0], [1, 1, 0, 0]], dtype=torch.bool)
+    cap_mask = torch.tensor([[1, 1, 0], [1, 0, 1]], dtype=torch.bool)
+
+    *_, unified_mask = UnifiedPrepare()(
+        x,
+        x_rope,
+        x_rope,
+        cap,
+        cap_rope,
+        cap_rope,
+        [4, 2],
+        [3, 3],
+        x_mask,
+        cap_mask,
+    )
+
+    assert torch.equal(
+        unified_mask,
+        torch.tensor(
+            [[1, 0, 1, 0, 1, 1, 0], [1, 1, 1, 0, 1, 0, 0]],
+            dtype=torch.bool,
+        ),
+    )
+
+
 def test_validate_zimage_tp_constraints_tp2_ok():
     ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints(
         dim=3840,

tests/model_executor/models/ming_image/test_config_and_pipeline.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.ming_image.pipeline import (
+    MingImageDiffusionPipeline,
+    _validate_variant_config,
+)
+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.models.ming_flash_omni.ming_flash_omni_thinker import (
+    MingFlashOmniThinkerForConditionalGeneration,
+)
+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.pipeline import MING_IMAGE_PIPELINE
+from vllm_omni.model_executor.stage_input_processors.ming_image import thinker2image
+from vllm_omni.model_extras import (
+    get_extra_body_params,
+    should_init_extra_args_for_non_diffusion_stages,
+)
+from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_config_selects_qwen25_vision_tower():
+    config = BailingMM2Config(vision_config={"model_type": "qwen2_5_vit"}, llm_config={})
+
+    assert type(config.vision_config).__name__ == "Qwen2_5_VLVisionConfig"
+
+
+@pytest.mark.parametrize(
+    ("transformer_config", "layered"),
+    (
+        (
+            SimpleNamespace(
+                _class_name="DiffusionTransformer",
+                alignment_padding_mode="zero_masked",
+                multi_frame_output=False,
+            ),
+            False,
+        ),
+        (
+            SimpleNamespace(
+                _class_name="DiffusionTransformer",
+                alignment_padding_mode="learned",
+                multi_frame_output=True,
+            ),
+            True,
+        ),
+    ),
+)
+def test_variant_config_without_model_index(transformer_config, layered):
+    assert _validate_variant_config(None, transformer_config) is layered
+
+
+def test_variant_config_rejects_mixed_semantics():
+    with pytest.raises(ValueError, match="must use either"):
+        _validate_variant_config(
+            {"_class_name": "MingImageDiffusionPipeline"},
+            SimpleNamespace(
+                _class_name="DiffusionTransformer",
+                alignment_padding_mode="learned",
+                multi_frame_output=False,
+            ),
+        )
+
+
+def test_layer_latent_frames_flatten_once_in_frame_major_order():
+    latents = torch.arange(2 * 4 * 3 * 2 * 2).reshape(2, 4, 3, 2, 2)
+
+    pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    flat = pipeline._flatten_latent_frames(latents)
+
+    assert flat.shape == (6, 4, 2, 2)
+    assert torch.equal(flat[0], latents[0, :, 0])
+    assert torch.equal(flat[1], latents[1, :, 0])
+    assert torch.equal(flat[2], latents[0, :, 1])
+
+
+def test_layer_pipeline_allows_missing_reference_only_for_dummy_run():
+    pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    pipeline.is_layer_decomposition = True
+    pipeline._configure_output_frames(
+        reference=None,
+        num_layers=1,
+        is_dummy_run=True,
+    )
+    assert pipeline._num_frames_per_prompt == 2
+
+    with pytest.raises(ValueError, match="requires a reference image"):
+        pipeline._configure_output_frames(
+            reference=None,
+            num_layers=1,
+            is_dummy_run=False,
+        )
+
+
+def test_pipeline_rejects_multiple_outputs_per_prompt():
+    pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    request = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="test",
+                sampling_params=OmniDiffusionSamplingParams(num_outputs_per_prompt=2),
+                request_id="test-request",
+            )
+        ]
+    )
+
+    with pytest.raises(ValueError, match="num_outputs_per_prompt=1 only, got 2"):
+        pipeline.forward(request)
+
+
+def _source_output(prefix_len: int = 4):
+    prompt_ids = [11] * prefix_len + [157158] + [157157] * 256 + [157159]
+    length = len(prompt_ids)
+    multimodal = {
+        "final_hidden_states": torch.arange(length * 2048).view(length, 2048),
+    }
+    for layer in (5, 12, 20):
+        multimodal[f"hidden_states_{layer}"] = torch.full((length, 2048), layer)
+    return SimpleNamespace(
+        prompt_token_ids=prompt_ids,
+        outputs=[SimpleNamespace(multimodal_output=multimodal)],
+    )
+
+
+def test_bridge_extracts_query_and_direct_conditions():
+    result = thinker2image(
+        [_source_output()],
+        prompt={"multi_modal_data": {"img2img": "reference"}},
+        sampling_params=SimpleNamespace(extra_args={"num_layers": 3}),
+    )
+
+    extra = result[0]["extra"]
+    assert extra["query_hidden_states"].shape == (256, 2048)
+    assert extra["direct_hidden_states"].shape == (4, 6144)
+    assert extra["reference_image"] == "reference"
+    assert extra["num_layers"] == 3
+    assert set(extra["direct_hidden_states"][:, :2048].unique().tolist()) == {5}
+
+
+def test_bridge_rejects_text_negative_condition():
+    with pytest.raises(ValueError, match="does not accept negative_prompt"):
+        thinker2image(
+            [_source_output()],
+            prompt={"negative_prompt": "bad"},
+            sampling_params=SimpleNamespace(extra_args={}),
+        )
+
+
+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"]
+    assert MING_IMAGE_PIPELINE.diffusers_class_name == "MingImageDiffusionPipeline"
+    assert MING_IMAGE_PIPELINE.diffusers_class_aliases == ("MingImageLayeredDiffusionPipeline",)
+    assert MING_IMAGE_PIPELINE.stages[0].model_subdir == "mllm"
+    assert MING_IMAGE_PIPELINE.stages[0].model_path_resolver.endswith(".resolve_ming_image_model_root")
+    assert MING_IMAGE_PIPELINE.stages[1].model_arch == "MingImageDiffusionPipeline"
+    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 should_init_extra_args_for_non_diffusion_stages(class_name)
+
+
+def test_stage1_compile_uses_configured_dynamic_regional_cuda_graph(monkeypatch):
+    pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        diffusion_compile_granularity="regional",
+        diffusion_compile_dynamic=True,
+    )
+    pipeline.transformer = torch.nn.Identity()
+    captured = {}
+
+    def _regionally_compile(model, **kwargs):
+        captured.update(kwargs)
+        return model
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.ming_image.pipeline.regionally_compile",
+        _regionally_compile,
+    )
+
+    pipeline.setup_compile()
+
+    assert captured == {
+        "mode": "reduce-overhead",
+        "fullgraph": True,
+        "dynamic": True,
+    }
+    assert pipeline._uses_cudagraph_trees
+
+
+def test_explicit_pipeline_loads_component_config_without_model_index(monkeypatch):
+    transformer_config = {
+        "_class_name": "DiffusionTransformer",
+        "alignment_padding_mode": "zero_masked",
+        "multi_frame_output": False,
+    }
+
+    def _get_hf_file_to_dict(filename, model, revision=None):
+        del model, revision
+        return transformer_config if filename == "transformer/config.json" else None
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+        lambda *_args, **_kwargs: None,
+    )
+    monkeypatch.setattr(
+        "vllm.transformers_utils.config.get_hf_file_to_dict",
+        _get_hf_file_to_dict,
+    )
+
+    config = OmniDiffusionConfig(
+        model="org/model",
+        model_class_name="MingImageDiffusionPipeline",
+    )
+    config.enrich_config()
+
+    assert config.model_class_name == "MingImageDiffusionPipeline"
+    assert config.tf_model_config.alignment_padding_mode == "zero_masked"
+    assert config.tf_model_config.multi_frame_output is False
+
+
+def test_checkpoint_resolver_downloads_mllm_and_sibling_mlp(monkeypatch):
+    captured = {}
+
+    def _download(**kwargs):
+        captured.update(kwargs)
+        return "/cache/model"
+
+    monkeypatch.setattr(checkpoint, "download_weights_from_hf_specific", _download)
+
+    assert checkpoint.resolve_ming_image_model_root("org/model", "rev", None) == "/cache/model"
+    assert captured["allow_patterns"] == ["mllm/**", "mlp/**"]
+    assert captured["revision"] == "rev"
+    assert captured["require_all"] is True
+
+
+def test_compute_logits_accepts_vllm_v1_signature():
+    class _LanguageModel(torch.nn.Module):
+        def compute_logits(self, hidden_states, sampling_metadata):
+            assert sampling_metadata is None
+            return hidden_states + 1
+
+    model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.language_model = _LanguageModel()
+
+    hidden_states = torch.zeros(2, 3)
+    assert torch.equal(model.compute_logits(hidden_states), hidden_states + 1)
+
+
+def test_embedding_only_run_uses_empty_modality_masks():
+    model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    image_mask, audio_mask = model._compute_modality_masks(None, torch.empty(7, 2048))
+
+    assert image_mask.shape == (7,)
+    assert not image_mask.any()
+    assert torch.equal(audio_mask, image_mask)
+
+
+def test_mllm_reports_root_mlp_query_tokens_as_loaded(monkeypatch):
+    model = MingImageForConditionalGeneration.__new__(MingImageForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.query_tokens_dict = torch.nn.ParameterDict({"16x16": torch.nn.Parameter(torch.empty(1), requires_grad=False)})
+    monkeypatch.setattr(
+        MingFlashOmniThinkerForConditionalGeneration,
+        "load_weights",
+        lambda self, weights: {"language_model.weight"},
+    )
+
+    assert model.load_weights([]) == {
+        "language_model.weight",
+        "query_tokens_dict.16x16",
+    }

tests/model_tests/diffusion/test_alignment.py

@@ -53,6 +53,8 @@ EXCLUDED_MODELS = [
     "BooguImageTurboPipeline",
     "LancePipeline",
     "MingImagePipeline",
+    "MingImageDiffusionPipeline",
+    "MingImageLayeredDiffusionPipeline",
     "InternVLAA1Pipeline",
     "StableDiffusion3Pipeline",
     "HunyuanImage3ForCausalMM",

vllm_omni/config/pipeline_registry.py

@@ -86,6 +86,7 @@ from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     MING_FLASH_OMNI_TTS_PIPELINE,
 )
+from vllm_omni.model_executor.models.ming_image.pipeline import MING_IMAGE_PIPELINE
 from vllm_omni.model_executor.models.ming_tts.pipeline import (
     MING_TTS_MOE_PIPELINE,
     MING_TTS_PIPELINE,
@@ -187,6 +188,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE,
     "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
+    "ming_image": MING_IMAGE_PIPELINE,
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
     "minimax_h3_disaggregated": MINIMAX_H3_PIPELINE,
     "omnivoice": OMNIVOICE_PIPELINE,

vllm_omni/deploy/ming_image.yaml

@@ -0,0 +1,46 @@
+# Ming-Image Design and Design-Layer two-stage deployment.
+pipeline: ming_image
+async_chunk: false
+
+trust_remote_code: false
+distributed_executor_backend: mp
+dtype: bfloat16
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    enforce_eager: true
+    tensor_parallel_size: 1
+    devices: "0"
+    gpu_memory_utilization: 0.55
+    max_num_seqs: 1
+    output_connectors:
+      to_stage_1: shared_memory_connector
+    default_sampling_params:
+      max_tokens: 1
+      temperature: 0
+      seed: 42
+
+  - stage_id: 1
+    enforce_eager: false
+    diffusion_compile_granularity: regional
+    diffusion_compile_dynamic: true
+    tensor_parallel_size: 1
+    devices: "1"
+    max_num_seqs: 1
+    input_connectors:
+      from_stage_0: shared_memory_connector
+    default_sampling_params:
+      seed: 42
+
+connectors:
+  shared_memory_connector:
+    name: SharedMemoryConnector
+
+edges:
+  - from: 0
+    to: 1
+    window_size: -1
+    max_inflight: 1
+
+platforms: {}

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -444,6 +444,7 @@ def extract_zimage_context(
         cap_pos_ids,
         x_inner_pad_mask,
         cap_inner_pad_mask,
+        _cap_feats_2,
     ) = module.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
 
     # Process image patches through embedder and noise refiner

vllm_omni/diffusion/data.py

@@ -1487,6 +1487,29 @@ class OmniDiffusionConfig:
         self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
         self.supports_mixed_reference_inputs = metadata.supports_mixed_reference_inputs
 
+    def _load_component_transformer_config(self) -> bool:
+        """Load a registered pipeline's component-level DiT configuration."""
+        from vllm.transformers_utils.config import get_hf_file_to_dict
+
+        from vllm_omni.model_extras import get_transformer_config_subfolder
+
+        transformer_subfolder = get_transformer_config_subfolder(
+            self.model_class_name,
+            model=self.model,
+            revision=self.revision,
+        )
+        tf_config_dict = get_hf_file_to_dict(
+            f"{transformer_subfolder}/config.json",
+            self.model,
+            revision=self.revision,
+        )
+        if tf_config_dict is None:
+            tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
+        if tf_config_dict is None:
+            return False
+        self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
+        return True
+
     @staticmethod
     def _looks_like_lance_subfolder(model: str | None) -> bool:
         """Return True when ``--model`` points at a Lance per-component subfolder.
@@ -1553,23 +1576,7 @@ class OmniDiffusionConfig:
                             exc,
                         )
                 else:
-                    from vllm_omni.model_extras import get_transformer_config_subfolder
-
-                    transformer_subfolder = get_transformer_config_subfolder(
-                        self.model_class_name,
-                        model=self.model,
-                        revision=self.revision,
-                    )
-                    tf_config_dict = get_hf_file_to_dict(
-                        f"{transformer_subfolder}/config.json",
-                        self.model,
-                        revision=self.revision,
-                    )
-                    if tf_config_dict is None:
-                        tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
-                    if tf_config_dict is not None:
-                        self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
-                    else:
+                    if not self._load_component_transformer_config():
                         self.set_tf_model_config(TransformerConfig())
             else:
                 raise FileNotFoundError("Diffusers pipeline index not found")
@@ -1605,6 +1612,23 @@ class OmniDiffusionConfig:
                         self.set_tf_model_config(TransformerConfig())
                         self.update_multimodal_support()
                         return
+                    # An explicit topology or CLI override can select a native,
+                    # registered pipeline whose repository contains component
+                    # configs but no root HF config or Diffusers index. Trust
+                    # that explicit class selection only when its component
+                    # transformer config is present.
+                    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+                    if (
+                        self.model_class_name in DiffusionModelRegistry.get_supported_archs()
+                        and self._load_component_transformer_config()
+                    ):
+                        logger.info(
+                            "Using explicitly selected diffusion pipeline %r with component-level configuration.",
+                            self.model_class_name,
+                        )
+                        self.update_multimodal_support()
+                        return
                     raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}")
 
                 self.set_tf_model_config(TransformerConfig.from_dict(cfg))

vllm_omni/diffusion/forward_context.py

@@ -48,6 +48,8 @@ class ForwardContext:
     total_denoise_steps: int | None = None
     # Per-request reference latent for img2img DiT models (e.g. Ming)
     ref_latent: torch.Tensor | None = None
+    # Per-request projected direct-VLM condition (e.g., Ming-Image). For now for bsz 1.
+    direct_condition: torch.Tensor | None = None
     # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
 
     # Sequence Parallel padding support
@@ -377,3 +379,9 @@ def set_forward_context_ref_latent(ref_latent: torch.Tensor | None) -> None:
     """
     if _forward_context is not None:
         _forward_context.ref_latent = ref_latent
+
+
+def set_forward_context_direct_condition(direct_condition: torch.Tensor | None) -> None:
+    """Set the projected direct-VLM condition on the active context."""
+    if _forward_context is not None:
+        _forward_context.direct_condition = direct_condition

vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 #
 # Adapted from Ming repository (inclusionAI/Ming) — the ``get_condition_embeds
@@ -65,12 +66,16 @@ class MingConditionEncoder(nn.Module):
         thinker_hidden_size: int = 4096,
         device: torch.device | str | None = None,
         dtype: torch.dtype | None = None,
+        normalize_output: bool = True,
+        strict_loading: bool = False,
     ) -> None:
         super().__init__()
         self.config = image_gen_config
         self.thinker_hidden_size = thinker_hidden_size
         self._target_device = torch.device(device) if device is not None else None
         self._target_dtype = dtype
+        self.normalize_output = normalize_output
+        self.strict_loading = strict_loading
 
         # Populated lazily by ``load_from_checkpoint`` to keep this module
         # cheap to construct (useful for dummy-init paths and unit tests).
@@ -195,6 +200,8 @@ class MingConditionEncoder(nn.Module):
         condition path to be meaningful.
         """
         if not mlp_path.exists():
+            if self.strict_loading:
+                raise FileNotFoundError(f"Missing Ming MLP directory: {mlp_path}")
             logger.warning(
                 "[MingConditionEncoder] mlp/ subfolder missing at %s — proj/norm "
                 "will stay randomly initialized. EXPECT BAD IMAGES until this "
@@ -206,6 +213,8 @@ class MingConditionEncoder(nn.Module):
         try:
             from safetensors.torch import load_file  # type: ignore
         except ImportError:
+            if self.strict_loading:
+                raise
             logger.exception("[MingConditionEncoder] safetensors not installed")
             return
 
@@ -213,6 +222,8 @@ class MingConditionEncoder(nn.Module):
         if not candidates:
             candidates = sorted(mlp_path.glob("*.bin"))
         if not candidates:
+            if self.strict_loading:
+                raise FileNotFoundError(f"No Ming MLP weights under {mlp_path}")
             logger.warning("[MingConditionEncoder] no weight files under %s", mlp_path)
             return
 
@@ -258,6 +269,8 @@ class MingConditionEncoder(nn.Module):
         ok_ow = _copy(self.proj_out.weight, "proj_out.weight")
         ok_ob = _copy(self.proj_out.bias, "proj_out.bias")
         if not (ok_w and ok_b and ok_ow and ok_ob):
+            if self.strict_loading:
+                raise ValueError(f"Ming condition projection weights are incomplete or incompatible under {mlp_path}.")
             logger.error(
                 "[MingConditionEncoder] proj_in/proj_out NOT fully loaded; diffusion conditioning will be garbage."
             )
@@ -342,7 +355,8 @@ class MingConditionEncoder(nn.Module):
         # L2 normalize + ×1000 rescale.  Ming's ``get_condition_embeds_for_image_gen``
         # applies F.normalize; ``ZImageModel_withMLP`` then rescales by 1000
         # when ``text_encoder_norm=True``.  We fold both into one place.
-        cap_feats = torch.nn.functional.normalize(cap_feats, dim=-1)
+        if self.normalize_output:
+            cap_feats = torch.nn.functional.normalize(cap_feats, dim=-1)
         if self.config.text_encoder_norm:
             cap_feats = cap_feats * 1000.0
 

vllm_omni/diffusion/models/ming_image/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Ming-Image specific diffusion components."""

vllm_omni/diffusion/models/ming_image/component_config.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Validation for Ming-Image component configuration."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+
+def load_mlp_config(model_path: str | Path) -> dict[str, Any]:
+    path = Path(model_path) / "mlp" / "config.json"
+    if not path.is_file():
+        raise FileNotFoundError(f"Ming-Image checkpoint is missing MLP config: {path}")
+    data = json.loads(path.read_text())
+    required = {
+        "diffusion_c_input_dim",
+        "diffusion_inner_dim",
+        "img_gen_scales",
+        "selected_hidden_states_layers",
+        "use_identity_mlp",
+        "use_learnable_token_condition",
+        "use_vlm_directvlm_condition",
+    }
+    missing = sorted(required - data.keys())
+    if missing:
+        raise ValueError(f"Ming-Image MLP config is missing keys: {missing}")
+    if data["img_gen_scales"] != [16]:
+        raise ValueError(f"Only img_gen_scales=[16] is supported, got {data['img_gen_scales']!r}.")
+    if data["selected_hidden_states_layers"] != [5, 12, 20]:
+        raise ValueError("Ming-Image requires selected hidden layers [5, 12, 20].")
+    if not data["use_identity_mlp"]:
+        raise NotImplementedError("Ming-Image non-identity diffusion MLP is not supported.")
+    if not data["use_learnable_token_condition"] or not data["use_vlm_directvlm_condition"]:
+        raise ValueError("Ming-Image requires both query-token and direct-VLM conditioning.")
+    return data

vllm_omni/diffusion/models/ming_image/condition.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ming-Image query and direct-VLM condition projectors."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file
+
+from vllm_omni.diffusion.models.ming_flash_omni.condition_encoder import MingConditionEncoder
+from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig
+
+from .component_config import load_mlp_config
+
+
+class MingImageConditioning(nn.Module):
+    def __init__(
+        self,
+        model_path: str | Path,
+        *,
+        device: torch.device | str,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__()
+        model_path = Path(model_path)
+        mlp_config = load_mlp_config(model_path)
+        config = MingImageGenConfig(
+            diffusion_c_input_dim=int(mlp_config["diffusion_c_input_dim"]),
+            thinker_hidden_size=2048,
+            img_gen_scales=list(mlp_config["img_gen_scales"]),
+            text_encoder_norm=bool(mlp_config.get("text_encoder_norm", False)),
+        )
+        self.query_encoder = MingConditionEncoder(
+            config,
+            thinker_hidden_size=2048,
+            device=device,
+            dtype=dtype,
+            normalize_output=bool(mlp_config.get("connector_norm", False)),
+            strict_loading=True,
+        )
+        self.query_encoder.load_from_checkpoint(model_path)
+
+        selected_layers = mlp_config["selected_hidden_states_layers"]
+        direct_input_dim = 2048 * len(selected_layers)
+        direct_output_dim = int(mlp_config["diffusion_inner_dim"])
+        self.direct_projector = nn.Sequential(
+            nn.RMSNorm(direct_input_dim, eps=1e-5),
+            nn.Linear(direct_input_dim, direct_output_dim, bias=True),
+        ).to(device=device, dtype=dtype)
+        self._load_direct_weights(model_path / "mlp" / "model.safetensors")
+
+    def _load_direct_weights(self, path: Path) -> None:
+        if not path.is_file():
+            raise FileNotFoundError(f"Missing Ming-Image direct projector weights: {path}")
+        state = load_file(str(path))
+        mapping = {
+            "0.weight": "proj_directvlm.0.weight",
+            "1.weight": "proj_directvlm.1.weight",
+            "1.bias": "proj_directvlm.1.bias",
+        }
+        own_state = self.direct_projector.state_dict()
+        projected: dict[str, torch.Tensor] = {}
+        for target, source in mapping.items():
+            if source not in state:
+                raise ValueError(f"Missing Ming-Image MLP weight {source!r}.")
+            if own_state[target].shape != state[source].shape:
+                raise ValueError(
+                    f"Ming-Image MLP shape mismatch for {source}: "
+                    f"expected {tuple(own_state[target].shape)}, got {tuple(state[source].shape)}."
+                )
+            projected[target] = state[source]
+        self.direct_projector.load_state_dict(projected, strict=True)
+
+    def forward(
+        self,
+        query_hidden_states: torch.Tensor,
+        direct_hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return (
+            self.query_encoder(query_hidden_states),
+            self.direct_projector(direct_hidden_states),
+        )
+
+
+__all__ = ["MingImageConditioning"]

vllm_omni/diffusion/models/ming_image/pipeline.py

@@ -0,0 +1,425 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Diffusion pipeline for Ming-Image Design and Design-Layer checkpoints."""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Iterable
+from typing import Any, ClassVar
+
+import torch
+import torch.nn as nn
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import (
+    DistributedAutoencoderKLQwenImage,
+)
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import (
+    set_forward_context_direct_condition,
+    set_forward_context_ref_latent,
+)
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch
+from vllm_omni.diffusion.models.ming_image.condition import MingImageConditioning
+from vllm_omni.diffusion.models.ming_image.transformer import MingImageTransformer2DModel
+from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+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
+
+logger = logging.getLogger(__name__)
+
+_DESIGN_PIPELINE = "MingImageDiffusionPipeline"
+_LAYERED_PIPELINE = "MingImageLayeredDiffusionPipeline"
+_VENDOR_TRANSFORMER_CLASS = "DiffusionTransformer"
+_DIFFUSION_REQUIRED_PATTERNS = [
+    "model_index.json",
+    "scheduler/**",
+    "transformer/**",
+    "vae/**",
+    "mlp/**",
+    "connector/**",
+]
+
+
+def _validate_variant_config(
+    model_index: dict[str, Any] | None,
+    transformer_config: Any,
+) -> bool:
+    """Validate checkpoint-owned variant metadata and return layered mode."""
+    declared_transformer = getattr(transformer_config, "_class_name", None)
+    if declared_transformer != _VENDOR_TRANSFORMER_CLASS:
+        raise ValueError(
+            "Ming-Image transformer/config.json must preserve the vendor "
+            f"_class_name={_VENDOR_TRANSFORMER_CLASS!r}, got {declared_transformer!r}. "
+            "vLLM-Omni loads those weights through MingImageTransformer2DModel."
+        )
+
+    alignment_padding_mode = getattr(transformer_config, "alignment_padding_mode", None)
+    multi_frame_output = getattr(transformer_config, "multi_frame_output", None)
+    variant = (alignment_padding_mode, multi_frame_output)
+    if variant == ("zero_masked", False):
+        is_layer_decomposition = False
+    elif variant == ("learned", True):
+        is_layer_decomposition = True
+    else:
+        raise ValueError(
+            "Ming-Image transformer/config.json must use either "
+            "alignment_padding_mode='zero_masked' with multi_frame_output=False "
+            "or alignment_padding_mode='learned' with multi_frame_output=True; "
+            f"got alignment_padding_mode={alignment_padding_mode!r}, "
+            f"multi_frame_output={multi_frame_output!r}."
+        )
+
+    declared_pipeline = (model_index or {}).get("_class_name")
+    if declared_pipeline is not None:
+        if declared_pipeline not in {_DESIGN_PIPELINE, _LAYERED_PIPELINE}:
+            raise ValueError(
+                "Ming-Image model_index.json must declare "
+                f"{_DESIGN_PIPELINE!r} or {_LAYERED_PIPELINE!r}, got {declared_pipeline!r}."
+            )
+        expected_pipeline = _LAYERED_PIPELINE if is_layer_decomposition else _DESIGN_PIPELINE
+        if declared_pipeline != expected_pipeline:
+            raise ValueError(
+                f"Ming-Image model_index.json declares {declared_pipeline!r}, but "
+                f"transformer/config.json describes {expected_pipeline!r}."
+            )
+    return is_layer_decomposition
+
+
+class MingImageDiffusionPipeline(ZImagePipeline):
+    """Ming-Image component adapter around the canonical Z-Image loop."""
+
+    supports_request_batch = False
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["conditioning"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
+        del prefix
+        nn.Module.__init__(self)
+
+        model_path = od_config.model
+        if not os.path.exists(model_path):
+            model_path = download_weights_from_hf_specific(
+                model_name_or_path=model_path,
+                cache_dir=None,
+                allow_patterns=_DIFFUSION_REQUIRED_PATTERNS,
+                revision=od_config.revision,
+                require_all=True,
+            )
+        local_files_only = os.path.isdir(model_path)
+        dtype = od_config.dtype
+
+        self.od_config = od_config
+        self._execution_device = get_local_device()
+        self.device = self._execution_device
+        model_index = get_diffusion_model_index(model_path, revision=od_config.revision) or {}
+        transformer_config = od_config.tf_model_config
+        self.is_layer_decomposition = _validate_variant_config(
+            model_index,
+            transformer_config,
+        )
+
+        self.default_num_inference_steps = 12
+        self.default_guidance_scale = 2.0 if self.is_layer_decomposition else 1.0
+        self._num_frames_per_prompt = 1
+        self._pending_prompt_embeds: list[torch.Tensor] | None = None
+        self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None
+        self._uses_cudagraph_trees = False
+
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder="transformer",
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+        subfolders = ["scheduler", "transformer", "vae"]
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_path,
+            subfolder="scheduler",
+            local_files_only=local_files_only,
+        )
+
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLQwenImage.from_pretrained,
+            model_path,
+            subfolder="vae",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
+        ).to(self.device)
+        self.vae.eval()
+        vae_channels = int(getattr(self.vae.config, "input_channels", 3))
+        if vae_channels != 4:
+            raise ValueError(f"Ming-Image requires an RGBA VAE, got input_channels={vae_channels}.")
+
+        transformer_kwargs = get_transformer_config_kwargs(
+            od_config.tf_model_config,
+            MingImageTransformer2DModel,
+        )
+        self.transformer = MingImageTransformer2DModel(
+            quant_config=od_config.quantization_config,
+            **transformer_kwargs,
+        )
+        self.conditioning = MingImageConditioning(
+            model_path,
+            device=self.device,
+            dtype=dtype,
+        )
+        self.text_encoder = None
+        self.tokenizer = None
+
+        self.vae_scale_factor = 2 ** len(self.vae.config.temperal_downsample)
+        self.image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor * 2,
+            do_convert_rgb=False,
+        )
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def setup_compile(self) -> None:
+        # Keep request preparation, scheduling, and VAE work eager,
+        # while capturing repeated DiT blocks for CUDAGraph Trees replay.
+        if self.od_config.diffusion_compile_granularity != "regional":
+            logger.warning(
+                "Ming-Image CUDA Graph uses regional DiT compilation; diffusion_compile_granularity=%r is ignored.",
+                self.od_config.diffusion_compile_granularity,
+            )
+        self.transformer = regionally_compile(
+            self.transformer,
+            mode="reduce-overhead",
+            fullgraph=True,
+            dynamic=self.od_config.diffusion_compile_dynamic,
+        )
+        self._uses_cudagraph_trees = True
+
+    def encode_prompt(self, *args, **kwargs):
+        del args, kwargs
+        if self._pending_prompt_embeds is None or self._pending_negative_prompt_embeds is None:
+            raise RuntimeError("Ming-Image conditioning is only available during forward.")
+        return self._pending_prompt_embeds, self._pending_negative_prompt_embeds
+
+    def prepare_latents(self, batch_size, *args, **kwargs):
+        frames = self._num_frames_per_prompt
+        flat = super().prepare_latents(batch_size * frames, *args, **kwargs)
+        return torch.stack(flat.chunk(frames, dim=0), dim=2)
+
+    def _encode_reference(
+        self,
+        reference: Any | None,
+        height: int,
+        width: int,
+    ) -> torch.Tensor | None:
+        if reference is None:
+            return None
+        if isinstance(reference, list):
+            if len(reference) != 1:
+                raise ValueError("Ming-Image currently accepts exactly one reference image.")
+            reference = reference[0]
+        if isinstance(reference, Image.Image):
+            reference = reference.convert("RGBA")
+        if not isinstance(reference, torch.Tensor):
+            reference = self.image_processor.preprocess(reference, height=height, width=width)
+        if reference.ndim != 4 or reference.shape[1] != 4:
+            raise ValueError("Ming-Image reference must be a batched RGBA tensor or an RGBA-compatible image.")
+        reference = reference.to(device=self.device, dtype=self.vae.dtype).unsqueeze(2)
+        latent = self.vae.encode(reference).latent_dist.mode()
+        return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+    @staticmethod
+    def _get_prompt_extra(req: DiffusionRequestBatch) -> dict[str, Any]:
+        if not req.prompts:
+            return {}
+        prompt = req.prompts[0]
+        if isinstance(prompt, dict):
+            return dict(prompt.get("extra") or {})
+        if hasattr(prompt, "_asdict"):
+            return dict(prompt._asdict().get("extra") or {})
+        return {}
+
+    def _configure_output_frames(
+        self,
+        *,
+        reference: Any | None,
+        num_layers: int,
+        is_dummy_run: bool,
+    ) -> None:
+        if num_layers < 1:
+            raise ValueError("num_layers must be at least 1.")
+        if self.is_layer_decomposition:
+            if reference is None and not is_dummy_run:
+                raise ValueError("Ming-Image Design-Layer requires a reference image.")
+            self._num_frames_per_prompt = num_layers + 1
+            return
+        if num_layers != 1:
+            raise ValueError("Ming-Image Design supports exactly one output frame.")
+        self._num_frames_per_prompt = 1
+
+    @torch.inference_mode()
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        sampling = req.sampling_params
+        if sampling.num_outputs_per_prompt != 1:
+            # TODO(yuanheng-zhao): enable after supporting batching
+            raise ValueError(
+                f"Ming-Image currently supports num_outputs_per_prompt=1 only, got {sampling.num_outputs_per_prompt}."
+            )
+
+        extra = self._get_prompt_extra(req)
+        extra_args = sampling.extra_args or {}
+
+        is_dummy_run = req.is_dummy_run()
+        query_hidden = extra.get("query_hidden_states")
+        direct_hidden = extra.get("direct_hidden_states")
+        if query_hidden is None or direct_hidden is None:
+            if not is_dummy_run:
+                raise ValueError("Ming-Image requests require query and direct conditions.")
+            logger.warning("Ming-Image conditions are absent during warmup; using zero tensors.")
+            query_hidden = torch.zeros((256, 2048), device=self.device, dtype=self.od_config.dtype)
+            direct_hidden = torch.zeros((1, 6144), device=self.device, dtype=self.od_config.dtype)
+        if not isinstance(query_hidden, torch.Tensor) or not isinstance(direct_hidden, torch.Tensor):
+            raise TypeError("Ming-Image query and direct conditions must be tensors.")
+        if query_hidden.ndim == 2:
+            query_hidden = query_hidden.unsqueeze(0)
+        if direct_hidden.ndim == 2:
+            direct_hidden = direct_hidden.unsqueeze(0)
+        query_hidden = query_hidden.to(device=self.device, dtype=self.od_config.dtype)
+        direct_hidden = direct_hidden.to(device=self.device, dtype=self.od_config.dtype)
+        cap_feats, direct_condition = self.conditioning(query_hidden, direct_hidden)
+
+        reference = extra.get("reference_image")
+        num_layers = int(extra_args.get("num_layers", extra.get("num_layers", 1)))
+        self._configure_output_frames(
+            reference=reference,
+            num_layers=num_layers,
+            is_dummy_run=is_dummy_run,
+        )
+
+        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
+        )
+        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
+        )
+
+        ref_latent = self._encode_reference(reference, height, width)
+        positive = [item for item in cap_feats]
+        negative = [torch.zeros_like(item) for item in positive]
+        self._pending_prompt_embeds = positive
+        self._pending_negative_prompt_embeds = negative
+
+        apply_cfg = cfg > 0
+        context_direct = (
+            torch.cat([direct_condition, torch.zeros_like(direct_condition)], dim=0) if apply_cfg else direct_condition
+        )
+        context_ref = ref_latent
+        if apply_cfg and context_ref is not None:
+            context_ref = context_ref.repeat(2, 1, 1, 1, 1)
+
+        inner_sampling = OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=cfg,
+            generator=generator,
+            output_type="latent",
+        )
+        inner_req = DiffusionRequestBatch(
+            requests=[
+                OmniDiffusionRequest(
+                    prompt={"prompt": ""},
+                    sampling_params=inner_sampling,
+                    request_id=req.request_id or "ming-image",
+                )
+            ]
+        )
+
+        set_forward_context_ref_latent(context_ref)
+        set_forward_context_direct_condition(context_direct)
+        try:
+            latent_output = super().forward(inner_req)
+            if not isinstance(latent_output.output, torch.Tensor):
+                raise TypeError("Ming-Image denoising must return latent tensors.")
+            image = self._decode_latent_frames(latent_output.output)
+            return DiffusionOutput(
+                output=image,
+                stage_durations=latent_output.stage_durations,
+            )
+        finally:
+            set_forward_context_ref_latent(None)
+            set_forward_context_direct_condition(None)
+            self._pending_prompt_embeds = None
+            self._pending_negative_prompt_embeds = None
+            self._num_frames_per_prompt = 1
+
+    @staticmethod
+    def _flatten_latent_frames(latents: torch.Tensor) -> torch.Tensor:
+        if latents.ndim == 4:
+            return latents
+        if latents.ndim != 5:
+            raise ValueError(f"Expected 4D or 5D Ming-Image latents, got {tuple(latents.shape)}")
+        batch, channels, frames, height, width = latents.shape
+        return latents.permute(2, 0, 1, 3, 4).reshape(
+            frames * batch,
+            channels,
+            height,
+            width,
+        )
+
+    def _decode_latent_frames(self, latents: torch.Tensor) -> torch.Tensor:
+        latents = self._flatten_latent_frames(latents).to(self.vae.dtype).unsqueeze(2)
+        latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents, return_dict=False)[0]
+        if image.ndim == 5:
+            if image.shape[2] != 1:
+                raise ValueError(
+                    f"Ming-Image VAE returned multiple decoded frames for a single latent frame: {tuple(image.shape)}"
+                )
+            image = image.squeeze(2)
+        return image
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded = AutoWeightsLoader(self).load_weights(weights)
+        loaded |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+        loaded |= {f"conditioning.{name}" for name, _ in self.conditioning.named_parameters()}
+        return loaded
+
+
+def get_ming_image_post_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+    image_processor = VaeImageProcessor(vae_scale_factor=16, do_convert_rgb=False)
+
+    def post_process(images: torch.Tensor):
+        if images.ndim == 5:
+            images = images.permute(0, 2, 1, 3, 4).flatten(0, 1)
+        return image_processor.postprocess(images.float())
+
+    return post_process
+
+
+__all__ = ["MingImageDiffusionPipeline", "get_ming_image_post_process_func"]

vllm_omni/diffusion/models/ming_image/transformer.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ming-Image adapter for the shared Z-Image transformer."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.models.z_image.z_image_transformer import ZImageTransformer2DModel
+
+
+class MingImageTransformer2DModel(ZImageTransformer2DModel):
+    """Runtime adapter for vendor DiffusionTransformer checkpoints.
+
+    This adapter reuses the shared Z-Image DiT and injects Ming-specific request conditions.
+    """
+
+    def forward(
+        self,
+        x: list[torch.Tensor],
+        t,
+        cap_feats: list[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+    ):
+        generated_frames = [item.shape[1] for item in x]
+        ref_x = None
+        cap_feats_2 = None
+        if is_forward_context_available():
+            context = get_forward_context()
+            if context.ref_latent is not None:
+                ref_x = [item.to(device=x[0].device, dtype=x[0].dtype) for item in context.ref_latent.unbind(dim=0)]
+            if context.direct_condition is not None:
+                cap_feats_2 = [
+                    item.to(device=x[0].device, dtype=x[0].dtype) for item in context.direct_condition.unbind(dim=0)
+                ]
+
+        output, metadata = super().forward(
+            x,
+            t,
+            cap_feats,
+            patch_size=patch_size,
+            f_patch_size=f_patch_size,
+            ref_x=ref_x,
+            cap_feats_2=cap_feats_2,
+        )
+        output = [item[:, :frames] for item, frames in zip(output, generated_frames)]
+        return output, metadata
+
+
+__all__ = ["MingImageTransformer2DModel"]

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
 #
@@ -584,7 +584,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
 
         # 5. Prepare timesteps
         if image is None:
-            image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2)
+            # for both [B, C, H, W] and multi-layer/frame [B, C, F, H, W]
+            image_seq_len = (latents.shape[-2] // 2) * (latents.shape[-1] // 2)
             mu = calculate_shift(
                 image_seq_len,
                 self.scheduler.config.get("base_image_seq_len", 256),
@@ -643,7 +644,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
             latents_typed = latents.to(self.od_config.dtype)
 
             if apply_cfg:
-                latent_model_input = latents_typed.repeat(2, 1, 1, 1)
+                repeat_dims = (2,) + (1,) * (latents_typed.ndim - 1)
+                latent_model_input = latents_typed.repeat(*repeat_dims)
                 prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds
                 timestep_model_input = timestep.repeat(2)
             else:
@@ -651,9 +653,13 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
                 prompt_embeds_model_input = prompt_embeds
                 timestep_model_input = timestep
 
-            latent_model_input = latent_model_input.unsqueeze(2)
+            if latent_model_input.ndim == 4:
+                latent_model_input = latent_model_input.unsqueeze(2)
             latent_model_input_list = list(latent_model_input.unbind(dim=0))
 
+            if getattr(self, "_uses_cudagraph_trees", False):
+                # check Ming-Image
+                torch.compiler.cudagraph_mark_step_begin()
             model_out_list = self.transformer(
                 latent_model_input_list,
                 timestep_model_input,
@@ -690,7 +696,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
             else:
                 noise_pred = torch.stack([t.float() for t in model_out_list], dim=0)
 
-            noise_pred = noise_pred.squeeze(2)
+            if latents.ndim == 4:
+                noise_pred = noise_pred.squeeze(2)
             noise_pred = -noise_pred
 
             # compute the previous noisy sample x_t -> x_t-1

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # _sp_plan definition adapted from HuggingFace diffusers library (_cp_plan)
 
 # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
@@ -56,6 +56,8 @@ from vllm_omni.model_executor.layers.timestep_embedding import timestep_embeddin
 ADALN_EMBED_DIM = 256
 SEQ_MULTI_OF = 32
 
+LEARNED_PADDING = "learned"
+ZERO_MASKED_PADDING = "zero_masked"
 logger = init_logger(__name__)
 
 
@@ -81,6 +83,8 @@ class UnifiedPrepare(nn.Module):
         cap_sin: torch.Tensor,
         x_item_seqlens: list[int],
         cap_item_seqlens: list[int],
+        x_attn_mask: torch.Tensor,
+        cap_attn_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """Combine x and cap tensors into unified sequences.
 
@@ -110,8 +114,9 @@ class UnifiedPrepare(nn.Module):
         unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0)
         unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0)
         unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
-        for i, seq_len in enumerate(unified_item_seqlens):
-            unified_attn_mask[i, :seq_len] = 1
+        for i, (x_len, cap_len) in enumerate(zip(x_item_seqlens, cap_item_seqlens)):
+            unified_attn_mask[i, :x_len] = x_attn_mask[i, :x_len]
+            unified_attn_mask[i, x_len : x_len + cap_len] = cap_attn_mask[i, :cap_len]
 
         return unified, unified_cos, unified_sin, unified_attn_mask
 
@@ -671,6 +676,8 @@ class ZImageTransformer2DModel(CachedTransformer):
         axes_dims=[32, 48, 48],
         axes_lens=[1024, 512, 512],
         quant_config: "QuantizationConfig | None" = None,
+        alignment_padding_mode: str = LEARNED_PADDING,
+        multi_frame_output: bool = True,
     ) -> None:
         super().__init__()
         # NOTE: `DiffusersPipelineLoader.load_model()` initializes this module
@@ -685,6 +692,12 @@ class ZImageTransformer2DModel(CachedTransformer):
         self.dim = dim
         self.n_heads = n_heads
 
+        if alignment_padding_mode not in {LEARNED_PADDING, ZERO_MASKED_PADDING}:
+            raise ValueError(f"Unsupported Z-Image alignment padding mode: {alignment_padding_mode!r}")
+        if type(multi_frame_output) is not bool:
+            raise ValueError("multi_frame_output must be boolean.")
+        self.alignment_padding_mode = alignment_padding_mode
+        self.multi_frame_output = multi_frame_output
         self.rope_theta = rope_theta
         self.t_scale = t_scale
         self.gradient_checkpointing = False
@@ -780,8 +793,12 @@ class ZImageTransformer2DModel(CachedTransformer):
             ),
         )
 
-        self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
-        self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
+        if alignment_padding_mode == LEARNED_PADDING:
+            self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
+            self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
+        else:
+            self.register_parameter("x_pad_token", None)
+            self.register_parameter("cap_pad_token", None)
 
         self.layers = nn.ModuleList(
             [
@@ -823,6 +840,8 @@ class ZImageTransformer2DModel(CachedTransformer):
                 .permute(6, 0, 3, 1, 4, 2, 5)
                 .reshape(self.out_channels, F, H, W)
             )
+            if not self.multi_frame_output:
+                x[i] = x[i][:, :1]
         return x
 
     @staticmethod
@@ -840,6 +859,8 @@ class ZImageTransformer2DModel(CachedTransformer):
         all_cap_feats: list[torch.Tensor],
         patch_size: int,
         f_patch_size: int,
+        all_image_ref: list[torch.Tensor | None] | None = None,
+        all_cap_feats_2: list[torch.Tensor | None] | None = None,
     ):
         pH = pW = patch_size
         pF = f_patch_size
@@ -852,10 +873,16 @@ class ZImageTransformer2DModel(CachedTransformer):
         all_cap_pos_ids = []
         all_cap_pad_mask = []
         all_cap_feats_out = []
+        all_cap_feats_2_out = []
 
-        for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)):
+        if all_image_ref is None:
+            all_image_ref = [None] * len(all_image)
+        if all_cap_feats_2 is None:
+            all_cap_feats_2 = [None] * len(all_image)
+
+        for image, cap_feat, cap_feat_2, image_ref in zip(all_image, all_cap_feats, all_cap_feats_2, all_image_ref):
             ### Process Caption
-            cap_ori_len = len(cap_feat)
+            cap_ori_len = len(cap_feat) + (len(cap_feat_2) if cap_feat_2 is not None else 0)
             cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
             # padded position ids
             cap_padded_pos_ids = self.create_coordinate_grid(
@@ -875,13 +902,17 @@ class ZImageTransformer2DModel(CachedTransformer):
                 )
             )
             # padded feature
-            cap_padded_feat = torch.cat(
-                [cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)],
-                dim=0,
-            )
-            all_cap_feats_out.append(cap_padded_feat)
+            if cap_feat_2 is None:
+                all_cap_feats_out.append(torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0))
+            else:
+                all_cap_feats_out.append(cap_feat)
+                all_cap_feats_2_out.append(
+                    torch.cat([cap_feat_2, torch.zeros_like(cap_feat_2[-1:]).repeat(cap_padding_len, 1)])
+                )
 
             ### Process Image
+            if image_ref is not None:
+                image = torch.cat([image, image_ref], dim=1)
             C, F, H, W = image.size()
             all_image_size.append((F, H, W))
             F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
@@ -931,6 +962,7 @@ class ZImageTransformer2DModel(CachedTransformer):
             all_cap_pos_ids,
             all_image_pad_mask,
             all_cap_pad_mask,
+            all_cap_feats_2_out,
         )
 
     def forward(
@@ -940,6 +972,8 @@ class ZImageTransformer2DModel(CachedTransformer):
         cap_feats: list[torch.Tensor],
         patch_size=2,
         f_patch_size=1,
+        ref_x: list[torch.Tensor | None] | None = None,
+        cap_feats_2: list[torch.Tensor] | None = None,
     ):
         assert patch_size in self.all_patch_size
         assert f_patch_size in self.all_f_patch_size
@@ -957,7 +991,8 @@ class ZImageTransformer2DModel(CachedTransformer):
             cap_pos_ids,
             x_inner_pad_mask,
             cap_inner_pad_mask,
-        ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
+            cap_feats_2,
+        ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size, ref_x, cap_feats_2)
 
         # x embed & refine
         x_item_seqlens = [len(_) for _ in x]
@@ -971,9 +1006,14 @@ class ZImageTransformer2DModel(CachedTransformer):
         adaln_input = t.type_as(x)
         # Use torch.where instead of x[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
         x_pad_mask = torch.cat(x_inner_pad_mask)
+        x_padding = (
+            self.x_pad_token.expand(x.shape[0], -1)
+            if self.alignment_padding_mode == LEARNED_PADDING
+            else torch.zeros_like(x)
+        )
         x = torch.where(
             x_pad_mask.unsqueeze(1).expand_as(x),
-            self.x_pad_token.expand(x.shape[0], -1),
+            x_padding,
             x,
         )
         x = list(x.split(x_item_seqlens, dim=0))
@@ -987,22 +1027,36 @@ class ZImageTransformer2DModel(CachedTransformer):
         x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device)
         for i, seq_len in enumerate(x_item_seqlens):
             x_attn_mask[i, :seq_len] = 1
+            if self.alignment_padding_mode == ZERO_MASKED_PADDING:
+                x_attn_mask[i, :seq_len].masked_fill_(x_inner_pad_mask[i], False)
 
         for layer in self.noise_refiner:
             x = layer(x, x_attn_mask, x_cos, x_sin, adaln_input)
 
         # cap embed & refine
         cap_item_seqlens = [len(_) for _ in cap_feats]
+        cap_feats = torch.cat(cap_feats, dim=0)
+        cap_feats = self.cap_embedder(cap_feats)
+        if cap_feats_2:
+            cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
+            if len(cap_feats) != len(cap_feats_2):
+                raise ValueError("Primary and direct caption conditions must have equal batch size.")
+            cap_feats = [torch.cat([primary, direct], dim=0) for primary, direct in zip(cap_feats, cap_feats_2)]
+            cap_item_seqlens = [len(item) for item in cap_feats]
+            cap_feats = torch.cat(cap_feats, dim=0)
         assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens)
         cap_max_item_seqlen = max(cap_item_seqlens)
 
-        cap_feats = torch.cat(cap_feats, dim=0)
-        cap_feats = self.cap_embedder(cap_feats)
         # Use torch.where instead of cap_feats[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
         cap_pad_mask = torch.cat(cap_inner_pad_mask)
+        cap_padding = (
+            self.cap_pad_token.expand(cap_feats.shape[0], -1)
+            if self.alignment_padding_mode == LEARNED_PADDING
+            else torch.zeros_like(cap_feats)
+        )
         cap_feats = torch.where(
             cap_pad_mask.unsqueeze(1).expand_as(cap_feats),
-            self.cap_pad_token.expand(cap_feats.shape[0], -1),
+            cap_padding,
             cap_feats,
         )
         cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
@@ -1016,6 +1070,8 @@ class ZImageTransformer2DModel(CachedTransformer):
         cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
         for i, seq_len in enumerate(cap_item_seqlens):
             cap_attn_mask[i, :seq_len] = 1
+            if self.alignment_padding_mode == ZERO_MASKED_PADDING:
+                cap_attn_mask[i, :seq_len].masked_fill_(cap_inner_pad_mask[i], False)
 
         for layer in self.context_refiner:
             cap_feats = layer(cap_feats, cap_attn_mask, cap_cos, cap_sin)
@@ -1023,9 +1079,17 @@ class ZImageTransformer2DModel(CachedTransformer):
         # Prepare unified tensors via UnifiedPrepare module
         # This enables _cp_plan to shard outputs via split_output=True
         unified, unified_cos, unified_sin, unified_attn_mask = self.unified_prepare(
-            x, x_cos, x_sin, cap_feats, cap_cos, cap_sin, x_item_seqlens, cap_item_seqlens
+            x,
+            x_cos,
+            x_sin,
+            cap_feats,
+            cap_cos,
+            cap_sin,
+            x_item_seqlens,
+            cap_item_seqlens,
+            x_attn_mask,
+            cap_attn_mask,
         )
-
         # Main transformer blocks
         for layer in self.layers:
             unified = layer(unified, unified_attn_mask, unified_cos, unified_sin, adaln_input)

vllm_omni/diffusion/registry.py

@@ -204,6 +204,11 @@ _DIFFUSION_MODELS = {
         "pipeline_ming_imagegen",
         "MingImagePipeline",
     ),
+    "MingImageDiffusionPipeline": (
+        "ming_image",
+        "pipeline",
+        "MingImageDiffusionPipeline",
+    ),
     "SanaWmPipeline": (
         "sana_wm",
         "pipeline_sana_wm",
@@ -390,6 +395,7 @@ _DIFFUSION_MODELS = {
         "Krea2Pipeline",
     ),
 }
+_DIFFUSION_MODELS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_MODELS["MingImageDiffusionPipeline"]
 
 
 DiffusionModelRegistry = _ModelRegistry(
@@ -663,6 +669,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "BagelPipeline": "get_bagel_post_process_func",
     "LancePipeline": "get_lance_post_process_func",
     "MingImagePipeline": "get_ming_image_post_process_func",
+    "MingImageDiffusionPipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "Pi0Pipeline": "get_pi0_post_process_func",
     "Pi05Pipeline": "get_pi05_post_process_func",
@@ -696,6 +703,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Krea2Pipeline": "get_krea2_post_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func",
 }
+_DIFFUSION_POST_PROCESS_FUNCS["MingImageLayeredDiffusionPipeline"] = _DIFFUSION_POST_PROCESS_FUNCS[
+    "MingImageDiffusionPipeline"
+]
 
 _DIFFUSION_IR_OP_PRIORITY_FUNCS = {
     # arch: ir_op_priority_func

vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 # Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved.
 # Adapted from Ming
@@ -728,7 +729,8 @@ class BailingMoeV2Model(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         image_mask: torch.Tensor | None = None,
         audio_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+        capture_layers: tuple[int, ...] | None = None,
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, dict[int, torch.Tensor]]:
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
@@ -740,7 +742,14 @@ class BailingMoeV2Model(nn.Module):
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
-        for layer in self.layers[self.start_layer : self.end_layer]:
+        capture_set = set(capture_layers or ())
+        captured: dict[int, torch.Tensor] = {}
+        for layer_idx, layer in enumerate(
+            self.layers[self.start_layer : self.end_layer],
+            start=self.start_layer,
+        ):
+            if layer_idx in capture_set:
+                captured[layer_idx] = hidden_states if residual is None else hidden_states + residual
             hidden_states, residual = layer(
                 positions,
                 hidden_states,
@@ -753,6 +762,10 @@ class BailingMoeV2Model(nn.Module):
             return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
 
         hidden_states, _ = self.norm(hidden_states, residual)
+        if self.end_layer in capture_set:
+            captured[self.end_layer] = hidden_states
+        if capture_layers:
+            return hidden_states, captured
         return hidden_states
 
 
@@ -811,6 +824,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin):
         inputs_embeds: torch.Tensor | None = None,
         image_mask: torch.Tensor | None = None,
         audio_mask: torch.Tensor | None = None,
+        capture_layers: tuple[int, ...] | None = None,
     ):
         hidden_states = self.model(
             input_ids=input_ids,
@@ -819,6 +833,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin):
             inputs_embeds=inputs_embeds,
             image_mask=image_mask,
             audio_mask=audio_mask,
+            capture_layers=capture_layers,
         )
         return hidden_states
 

vllm_omni/model_executor/models/ming_image/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

vllm_omni/model_executor/models/ming_image/checkpoint.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Checkpoint resolution for the Ming-Image MLLM stage."""
+
+from pathlib import Path
+
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+
+_MLLM_REQUIRED_PATTERNS = [
+    "mllm/**",
+    "mlp/**",
+]
+
+
+def resolve_ming_image_model_root(
+    model: str,
+    revision: str | None,
+    task_type: str | None,
+) -> str:
+    """Materialize the MLLM and its sibling query-token checkpoint.
+
+    The autoregressive weights and tokenizer live under mllm/,
+    while the learned image-query tokens consumed by that stage live under mlp/.
+    Stage subdirectory resolution runs after this hook, so return the repository
+    root and let model_subdir/tokenizer_subdir select mllm/.
+    """
+    del task_type
+    path = Path(model)
+    if path.is_dir():
+        if path.name == "mllm" and (path / "config.json").is_file():
+            return str(path.parent)
+        return str(path)
+
+    return download_weights_from_hf_specific(
+        model_name_or_path=model,
+        cache_dir=None,
+        allow_patterns=_MLLM_REQUIRED_PATTERNS,
+        revision=revision,
+        require_all=True,
+    )
+
+
+__all__ = ["resolve_ming_image_model_root"]

vllm_omni/model_executor/models/ming_image/model.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""vLLM-native MLLM stage for inclusionAI Ming-Image."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from vllm.config import VllmConfig
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
+from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.parse import MultiModalDataParser
+
+from vllm_omni.model_executor.models.ming_flash_omni.ming_flash_omni_thinker import (
+    MingFlashOmniThinkerDummyInputsBuilder,
+    MingFlashOmniThinkerForConditionalGeneration,
+    MingFlashOmniThinkerMultiModalProcessor,
+    MingFlashOmniThinkerProcessingInfo,
+)
+from vllm_omni.model_executor.models.ming_flash_omni.modeling_bailing_moe_v2 import (
+    BailingMoeV2ForCausalLM,
+)
+from vllm_omni.model_executor.models.ming_flash_omni.projectors import VisionProjector
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
+from vllm_omni.transformers_utils.processors.ming import MingImageProcessor
+
+
+class _MingImageDataParser(MultiModalDataParser):
+    def parse_mm_data(self, mm_data):
+        normalized = dict(mm_data)
+        if "img2img" in normalized:
+            references = normalized.pop("img2img")
+            references = references if isinstance(references, list) else [references]
+            existing = normalized.get("image", [])
+            existing = existing if isinstance(existing, list) else [existing]
+            normalized["image"] = existing + references
+        return super().parse_mm_data(normalized)
+
+
+class MingImageProcessingInfo(MingFlashOmniThinkerProcessingInfo):
+    def get_hf_config(self) -> BailingMM2Config:
+        return self.ctx.get_hf_config(BailingMM2Config)
+
+    def get_hf_processor(self, **kwargs: object):
+        return self.ctx.get_hf_processor(MingImageProcessor, **kwargs)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"image": None, "img2img": 1}
+
+    def get_mm_max_tokens_per_item(self, seq_len, mm_counts):
+        counts = dict(mm_counts or {})
+        counts["image"] = counts.get("image", 0) + counts.pop("img2img", 0)
+        return super(MingFlashOmniThinkerProcessingInfo, self).get_mm_max_tokens_per_item(
+            seq_len=seq_len,
+            mm_counts=counts,
+        )
+
+    def get_data_parser(self):
+        return _MingImageDataParser()
+
+
+class MingImageDummyInputsBuilder(MingFlashOmniThinkerDummyInputsBuilder):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        count = mm_counts.get("image", 0) + mm_counts.get("img2img", 0)
+        return self.info.get_hf_processor().image_token * count
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options=None,
+    ) -> MultiModalDataDict:
+        count = mm_counts.get("image", 0) + mm_counts.get("img2img", 0)
+        return {
+            "image": self._get_dummy_images(
+                width=448,
+                height=448,
+                num_images=count,
+            )
+        }
+
+
+class MingImageMultiModalProcessor(MingFlashOmniThinkerMultiModalProcessor):
+    pass
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MingImageMultiModalProcessor,
+    info=MingImageProcessingInfo,
+    dummy_inputs=MingImageDummyInputsBuilder,
+)
+class MingImageForConditionalGeneration(MingFlashOmniThinkerForConditionalGeneration):
+    """Ming-Image MLLM stage with Qwen2.5-VL and direct-state capture."""
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.": "language_model.",
+            "linear_proj.": "linear_proj.proj.",
+        },
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        torch.nn.Module.__init__(self)
+        thinker_config = vllm_config.model_config.hf_config
+        if not isinstance(thinker_config, BailingMM2Config):
+            raise TypeError(f"Expected BailingMM2Config, got {type(thinker_config).__name__}")
+        if thinker_config.llm_config is None or thinker_config.vision_config is None:
+            raise ValueError("Ming-Image requires llm_config and vision_config.")
+
+        llm_config = thinker_config.llm_config
+        self.config = llm_config
+        self.thinker_config = thinker_config
+        self.have_multimodal_outputs = True
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = BailingMoeV2ForCausalLM(
+                vllm_config=vllm_config.with_hf_config(llm_config),
+                prefix=maybe_prefix(prefix, "llm"),
+            )
+
+        with self._mark_tower_model(vllm_config, "image"):
+            self.vision = Qwen2_5_VisionTransformer(
+                thinker_config.vision_config,
+                norm_eps=llm_config.rms_norm_eps,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "vision"),
+            )
+            self.linear_proj = VisionProjector(
+                vision_dim=thinker_config.vision_config.out_hidden_size,
+                llm_dim=llm_config.hidden_size,
+                mlp_depth=getattr(thinker_config, "mlp_depth", 2),
+            )
+
+        self.audio = None
+        self.linear_proj_audio = None
+        self.query_tokens_dict = torch.nn.ParameterDict()
+        component_path = Path(vllm_config.model_config.model)
+        model_root = component_path.parent if component_path.name == "mllm" else component_path
+        self._load_image_gen_query_tokens(str(model_root))
+        if not self.query_tokens_dict:
+            raise ValueError(f"Ming-Image query tokens were not found under {model_root / 'mlp'}")
+
+        self.capture_layers = (5, 12, llm_config.num_hidden_layers)
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    def extract_image_feature(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
+        with torch.amp.autocast(pixel_values.device.type, dtype=torch.bfloat16):
+            image_embeds = self.vision(pixel_values, grid_thw=grid_thw.tolist())
+        return F.normalize(self.linear_proj(image_embeds), dim=-1)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata=None,
+    ) -> torch.Tensor | None:
+        # vLLM v1 calls this with hidden states only; the reused Bailing model
+        # still accepts the legacy metadata argument (which it does not use).
+        return self.language_model.compute_logits(hidden_states, sampling_metadata)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_weights = super().load_weights(weights)
+        # Query tokens live in the root-level mlp checkpoint and are loaded by
+        # ``_load_image_gen_query_tokens`` during construction, outside vLLM's
+        # main checkpoint iterator. Include them in the returned set so strict
+        # post-load validation accounts for that checkpoint-owned component.
+        loaded_weights.update(f"query_tokens_dict.{scale}" for scale in self.query_tokens_dict)
+        return loaded_weights
+
+    def _compute_modality_masks(
+        self,
+        input_ids: torch.Tensor | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_ids is None:
+            if inputs_embeds is None:
+                raise ValueError("input_ids and inputs_embeds cannot both be None")
+            empty_mask = torch.zeros(
+                inputs_embeds.shape[:-1],
+                device=inputs_embeds.device,
+                dtype=torch.bool,
+            )
+            return empty_mask, empty_mask
+        return input_ids == self.config.image_patch_token, torch.zeros_like(input_ids, dtype=torch.bool)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors=None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ) -> OmniOutput:
+        image_mask, audio_mask = self._compute_modality_masks(input_ids, inputs_embeds)
+        result = self.language_model.forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            image_mask=image_mask,
+            audio_mask=audio_mask,
+            capture_layers=self.capture_layers,
+        )
+        if not isinstance(result, tuple):
+            return OmniOutput(text_hidden_states=result, multimodal_outputs={"final_hidden_states": result})
+
+        hidden_states, captured = result
+        multimodal_outputs: dict[str, Any] = {"final_hidden_states": hidden_states}
+        for layer_idx, layer_hidden in captured.items():
+            multimodal_outputs[f"hidden_states_{layer_idx}"] = layer_hidden
+        return OmniOutput(
+            text_hidden_states=hidden_states,
+            multimodal_outputs=multimodal_outputs,
+        )
+
+
+__all__ = ["MingImageForConditionalGeneration"]

vllm_omni/model_executor/models/ming_image/pipeline.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Two-stage Ming-Image topology."""
+
+from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.ming_image"
+_CHECKPOINT = "vllm_omni.model_executor.models.ming_image.checkpoint"
+
+MING_IMAGE_PIPELINE = PipelineConfig(
+    model_type="ming_image",
+    default_deploy_config_name="ming_image.yaml",
+    model_arch="MingImageForConditionalGeneration",
+    hf_architectures=("MingImageForConditionalGeneration",),
+    diffusers_class_name="MingImageDiffusionPipeline",
+    diffusers_class_aliases=("MingImageLayeredDiffusionPipeline",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="mllm",
+            model_arch="MingImageForConditionalGeneration",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+            model_subdir="mllm",
+            tokenizer_subdir="mllm",
+            sampling_constraints={"detokenize": False},
+            model_path_resolver=f"{_CHECKPOINT}.resolve_ming_image_model_root",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            model_arch="MingImageDiffusionPipeline",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            custom_process_input_func=f"{_PROC}.thinker2image",
+        ),
+    ),
+)
+
+__all__ = ["MING_IMAGE_PIPELINE"]

vllm_omni/model_executor/models/registry.py

@@ -425,6 +425,11 @@ _OMNI_MODELS = {
         "ming_flash_omni_thinker",
         "MingFlashOmniThinkerForConditionalGeneration",
     ),
+    "MingImageForConditionalGeneration": (
+        "ming_image",
+        "model",
+        "MingImageForConditionalGeneration",
+    ),
     "MingFlashOmniTalkerForConditionalGeneration": (
         "ming_flash_omni",
         "ming_flash_omni_talker",

vllm_omni/model_executor/stage_input_processors/ming_image.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Strict MLLM-to-diffusion bridge for Ming-Image."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import _ensure_list
+
+_IMAGE_PATCH_TOKEN_ID = 157157
+_IMAGE_END_TOKEN_ID = 157159
+_NUM_QUERY_TOKENS = 256
+_CAPTURE_LAYERS = (5, 12, 20)
+
+
+def _output_payload(source_output: Any) -> tuple[list[int], dict[str, torch.Tensor]]:
+    if not getattr(source_output, "outputs", None):
+        raise ValueError("Ming-Image thinker produced no outputs.")
+    output = source_output.outputs[0]
+    multimodal = getattr(output, "multimodal_output", None) or {}
+    return _ensure_list(source_output.prompt_token_ids), multimodal
+
+
+def _validate_query_suffix(prompt_ids: list[int]) -> tuple[int, int]:
+    expected = _NUM_QUERY_TOKENS + 2
+    if len(prompt_ids) < expected:
+        raise ValueError("Ming-Image prompt is shorter than its image-query suffix.")
+    start = len(prompt_ids) - _NUM_QUERY_TOKENS - 1
+    end = len(prompt_ids) - 1
+    if prompt_ids[end] != _IMAGE_END_TOKEN_ID or any(token != _IMAGE_PATCH_TOKEN_ID for token in prompt_ids[start:end]):
+        raise ValueError("Ming-Image query suffix must end with 256 image-patch tokens and an image-end token.")
+    return start, end
+
+
+def _as_sequence(hidden: torch.Tensor, prompt_length: int, name: str) -> torch.Tensor:
+    if hidden.ndim == 3:
+        if hidden.shape[0] != 1:
+            raise ValueError(f"{name} must have batch size 1, got {tuple(hidden.shape)}")
+        hidden = hidden[0]
+    if hidden.ndim != 2 or hidden.shape[0] != prompt_length:
+        raise ValueError(f"{name} shape {tuple(hidden.shape)} does not match prompt length {prompt_length}.")
+    return hidden
+
+
+def _extract_conditions(source_output: Any) -> tuple[torch.Tensor, torch.Tensor]:
+    prompt_ids, multimodal = _output_payload(source_output)
+    query_start, query_end = _validate_query_suffix(prompt_ids)
+    final_hidden = multimodal.get("final_hidden_states")
+    if not isinstance(final_hidden, torch.Tensor):
+        raise ValueError("Ming-Image thinker output is missing final_hidden_states.")
+    final_hidden = _as_sequence(final_hidden, len(prompt_ids), "final_hidden_states")
+    query_hidden = final_hidden[query_start:query_end].detach().contiguous()
+
+    selected: list[torch.Tensor] = []
+    for layer_idx in _CAPTURE_LAYERS:
+        value = multimodal.get(f"hidden_states_{layer_idx}")
+        if not isinstance(value, torch.Tensor):
+            raise ValueError(f"Ming-Image thinker output is missing hidden_states_{layer_idx}.")
+        selected.append(_as_sequence(value, len(prompt_ids), f"hidden_states_{layer_idx}"))
+
+    # Vendor labels mark the full generated query block, including delimiters.
+    # Direct VLM conditioning is restricted to the original prompt/reference
+    # prefix and therefore excludes all 258 appended tokens.
+    direct_end = len(prompt_ids) - (_NUM_QUERY_TOKENS + 2)
+    direct_hidden = torch.cat([value[:direct_end] for value in selected], dim=-1)
+    return query_hidden, direct_hidden.detach().contiguous()
+
+
+def thinker2image(
+    source_outputs: list[Any],
+    prompt: Any | None = None,
+    requires_multimodal_data: bool = False,
+    sampling_params: Any | None = None,
+) -> list[dict[str, Any]]:
+    del requires_multimodal_data
+    if len(source_outputs) != 1:
+        raise ValueError(f"Ming-Image expects one thinker output, got {len(source_outputs)}.")
+
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    negative_prompt = extra_args.get("negative_prompt")
+    if isinstance(prompt, dict):
+        negative_prompt = prompt.get("negative_prompt", negative_prompt)
+    if isinstance(negative_prompt, str) and negative_prompt.strip():
+        raise ValueError("Ming-Image uses zero negative conditioning and does not accept negative_prompt.")
+
+    query_hidden, direct_hidden = _extract_conditions(source_outputs[0])
+    extra: dict[str, Any] = {
+        "query_hidden_states": query_hidden,
+        "direct_hidden_states": direct_hidden,
+    }
+
+    if isinstance(prompt, dict):
+        mm_data = prompt.get("multi_modal_data") or {}
+        reference = mm_data.get("img2img", mm_data.get("image"))
+        if isinstance(reference, list):
+            reference = reference[0] if reference else None
+        if reference is not None:
+            extra["reference_image"] = reference
+
+    num_layers = int(extra_args.get("num_layers", 1))
+    if num_layers < 1:
+        raise ValueError("num_layers must be at least 1.")
+    extra["num_layers"] = num_layers
+    return [{"prompt": "", "extra": extra}]
+
+
+__all__ = ["thinker2image"]

vllm_omni/model_extras/ming_image.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from vllm_omni.model_extras.ming_flash_omni import (
+    build_image_to_image_prompt,
+    build_text_to_image_prompt,
+)
+
+MING_IMAGE_EXTRA_BODY_PARAMS = frozenset(
+    {
+        "height",
+        "width",
+        "steps",
+        "cfg",
+        "seed",
+        "num_layers",
+        "negative_prompt",
+    }
+)
+MING_IMAGE_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset()
+MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True
+
+__all__ = [
+    "MING_IMAGE_EXTRA_BODY_PARAMS",
+    "MING_IMAGE_EXTRA_OUTPUT_PARAMS",
+    "MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES",
+    "build_image_to_image_prompt",
+    "build_text_to_image_prompt",
+]

vllm_omni/model_extras/registry.py

@@ -75,6 +75,17 @@ from vllm_omni.model_extras.ming_flash_omni import (
 from vllm_omni.model_extras.ming_flash_omni import (
     build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt,
 )
+from vllm_omni.model_extras.ming_image import (
+    MING_IMAGE_EXTRA_BODY_PARAMS,
+    MING_IMAGE_EXTRA_OUTPUT_PARAMS,
+    MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+)
+from vllm_omni.model_extras.ming_image import (
+    build_image_to_image_prompt as build_ming_image_image_to_image_prompt,
+)
+from vllm_omni.model_extras.ming_image import (
+    build_text_to_image_prompt as build_ming_image_text_to_image_prompt,
+)
 from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS
 from vllm_omni.model_extras.sensenova_u1 import (
     SENSENOVA_U1_EXTRA_BODY_PARAMS,
@@ -299,7 +310,15 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt,
         "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt,
     },
+    "MingImageDiffusionPipeline": {
+        "extra_body_params": MING_IMAGE_EXTRA_BODY_PARAMS,
+        "extra_output_params": MING_IMAGE_EXTRA_OUTPUT_PARAMS,
+        "init_extra_args_for_non_diffusion_stages": MING_IMAGE_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+        "text_to_image_prompt_builder": build_ming_image_text_to_image_prompt,
+        "image_to_image_prompt_builder": build_ming_image_image_to_image_prompt,
+    },
 }
+_EXTRA_SPECS["MingImageLayeredDiffusionPipeline"] = _EXTRA_SPECS["MingImageDiffusionPipeline"]
 
 for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"):
     _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder

vllm_omni/transformers_utils/configs/ming_flash_omni.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 # Copyright 2024 ANT Group and the HuggingFace Inc. team. All rights reserved.
 #
@@ -19,7 +20,13 @@
 import os
 from typing import Any, ClassVar
 
-from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerFast
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedTokenizerFast,
+    Qwen2_5_VLVisionConfig,
+)
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
@@ -251,9 +258,19 @@ class BailingMM2Config(PretrainedConfig):
         **kwargs,
     ):
         self.audio_config = WhisperEncoderConfig(**audio_config) if isinstance(audio_config, dict) else audio_config
-        self.vision_config = (
-            Qwen3VLMoeVisionConfig(**vision_config) if isinstance(vision_config, dict) else vision_config
-        )
+        if isinstance(vision_config, dict):
+            vision_config = dict(vision_config)
+            vision_model_type = vision_config.get("model_type")
+            if vision_model_type in {"qwen2_5_vit", "qwen2_5_vl"}:
+                # 1) Ming-Image uses the Qwen2.5-VL vision tower while
+                # 2) Ming-flash-omni-2.0 uses Qwen3-VL. Keep the composition config shared
+                # and select the concrete tower from ckpt metadata.
+                vision_config.pop("model_type", None)
+                self.vision_config = Qwen2_5_VLVisionConfig(**vision_config)
+            else:
+                self.vision_config = Qwen3VLMoeVisionConfig(**vision_config)
+        else:
+            self.vision_config = vision_config
         self.llm_config = BailingMoeV2Config(**llm_config) if isinstance(llm_config, dict) else llm_config
         self.mlp_depth = mlp_depth
         super().__init__(**kwargs)

vllm_omni/transformers_utils/processors/__init__.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 
 from vllm_omni.transformers_utils.processors.ming import (
     MingFlashOmniProcessor,
+    MingImageProcessor,
     MingWhisperFeatureExtractor,
 )
 
 __all__ = [
     "MingFlashOmniProcessor",
+    "MingImageProcessor",
     "MingWhisperFeatureExtractor",
 ]

vllm_omni/transformers_utils/processors/ming.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The vLLM-Omni team.
 # Copyright 2024 ANT Group and the HuggingFace Inc. team.
 #
@@ -18,7 +19,7 @@ from typing import Any
 
 import numpy as np
 import torch
-from transformers import AutoFeatureExtractor, AutoProcessor
+from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, Qwen2VLImageProcessor
 from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
 from transformers.processing_utils import ProcessorMixin
 from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
@@ -481,6 +482,87 @@ class MingFlashOmniProcessor(ProcessorMixin):
         return list(dict.fromkeys(names))
 
 
+class MingImageProcessor(ProcessorMixin):
+    """Image-only Bailing processor used by Ming-Image ckpts.
+
+    Equivalent to Qwen2-VL preprocessing. Construct the upstream component explicitly
+    in case model repo auto_map points at files that are not bundled in the ckpt.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "Qwen2VLImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, image_processor=None, tokenizer=None, merge_size: int = 2, **kwargs):
+        if image_processor is None or tokenizer is None:
+            raise ValueError("MingImageProcessor requires image_processor and tokenizer.")
+        self.spatial_merge_size = int(getattr(image_processor, "merge_size", merge_size))
+        self.image_token = PLACEHOLDER_IMAGE_TOKEN_IN_TEXT
+        super().__init__(image_processor=image_processor, tokenizer=tokenizer, **kwargs)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        image_processor = Qwen2VLImageProcessor.from_pretrained(
+            pretrained_model_name_or_path,
+            *args,
+            **kwargs,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            *args,
+            **kwargs,
+        )
+        return cls(image_processor=image_processor, tokenizer=tokenizer)
+
+    def __call__(self, text, images=None, **kwargs) -> BatchFeature:
+        if isinstance(text, str):
+            text = [text]
+        data: dict[str, Any] = {}
+        if images is not None:
+            image_outputs = self.image_processor(
+                images=images,
+                return_tensors="pt",
+                **kwargs.get("images_kwargs", {}),
+            )
+            data.update(image_outputs)
+            text = self._expand_image_tokens(text, image_outputs["image_grid_thw"])
+        data.update(
+            self.tokenizer(
+                text,
+                return_tensors="pt",
+                **kwargs.get("text_kwargs", {}),
+            )
+        )
+        return BatchFeature(data=data)
+
+    def _expand_image_tokens(self, text: list[str], image_grid_thw: torch.Tensor) -> list[str]:
+        num_patches = image_grid_thw.prod(dim=1) // (self.spatial_merge_size**2)
+        image_index = 0
+        expanded: list[str] = []
+        for sample in text:
+            for _ in range(sample.count(self.image_token)):
+                replacement = (
+                    DEFAULT_IM_START_TOKEN
+                    + DEFAULT_IMAGE_PATCH_TOKEN * int(num_patches[image_index].item())
+                    + DEFAULT_IM_END_TOKEN
+                    + "\n"
+                )
+                sample = sample.replace(self.image_token, replacement, 1)
+                image_index += 1
+            expanded.append(sample)
+        return expanded
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        return list(dict.fromkeys(self.tokenizer.model_input_names + self.image_processor.model_input_names))
+
+
 # transformers >= 5.x requires the config *class* (not a string) as the first
 # arg to AutoFeatureExtractor/AutoProcessor.register — it does `key.__module__`
 # internally. Passing "MingWhisperFeatureExtractor" raises: