[Model] Add Ming-flash-omni-2.0 Image Generation (Diffusion) Stage (#2875)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>

docs/models/supported_models.md

@@ -19,6 +19,7 @@ th {
 |--------------|--------|-------------------|------------|---------|-----|-----------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

examples/offline_inference/ming_flash_omni/README.md

@@ -2,12 +2,13 @@
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 
-vLLM-Omni supports two deployment modes:
+vLLM-Omni supports three deployment modes:
 
 | Mode | Deploy config | Output |
 |------|--------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -87,6 +88,17 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
 
 Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
 
+### Image generation (text-to-image / img2img)
+
+The diffusion-stage image-generation path is currently only wired through
+the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
+`"modalities": ["image"]`). For payloads, optional knobs
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
+reference-image flow, see the
+[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
+`end2end.py` does not yet exercise the imagegen stage; that is tracked as
+follow-up work.
+
 ## Online serving
 
 For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).

examples/online_serving/ming_flash_omni/README.md

@@ -10,6 +10,7 @@ Please refer to [README.md](../../../README.md)
 |------|---------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -56,8 +57,98 @@ bash run_curl_multimodal_generation.sh use_image
 bash run_curl_multimodal_generation.sh use_audio
 bash run_curl_multimodal_generation.sh use_video
 bash run_curl_multimodal_generation.sh use_mixed_modalities
+bash run_curl_multimodal_generation.sh use_image_gen
 ```
 
+## Image generation (text-to-image)
+
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
+Launch with the dual-stage YAML, which adds an image-gen stage behind the
+thinker:
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
+    --stage-init-timeout 1800 \
+    --port 8188
+```
+
+Then request image output by passing `"modalities": ["image"]`:
+
+```bash
+curl http://127.0.0.1:8188/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
+      "modalities": ["image"]
+    }' -o /tmp/ming_response.json
+
+python - <<'PY'
+import base64, json
+r = json.load(open('/tmp/ming_response.json'))
+url = r['choices'][0]['message']['content'][0]['image_url']['url']
+open('/tmp/ming_gen.png', 'wb').write(base64.b64decode(url.split(',', 1)[1]))
+print('saved /tmp/ming_gen.png')
+PY
+```
+
+### Optional knobs
+
+Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
+| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
+| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
+
+Example with all the knobs:
+
+```bash
+curl http://127.0.0.1:8188/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "modalities": ["image"],
+      "sampling_params_list": [
+        {"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
+         "extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
+        {"seed":42,"extra_args":{"image_gen":{
+            "steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
+            "byte5_text":["理解与生成统一"]}}}
+      ],
+      "messages": [{"role": "user", "content": "Draw a poster."}]
+    }' -o /tmp/ming_response.json
+```
+
+### img2img (reference image)
+
+Add an `image_url` content part to the user message; the parser routes it
+into the diffusion stage as `extra[reference_image]`:
+
+```jsonc
+"messages": [{
+  "role": "user",
+  "content": [
+    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
+    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
+  ]
+}]
+```
+
+### GPU layout
+
+The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
+(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
+`runtime.devices` per stage to relocate; with fewer GPUs available, drop
+the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
+warmup takes roughly an extra 30–60 s on top of the thinker — set
+`--stage-init-timeout 1800` if the default 300 s is too tight.
+
 ## Modality control
 
 | `modalities` | Server config | Output |
@@ -65,6 +156,7 @@ bash run_curl_multimodal_generation.sh use_mixed_modalities
 | `["text"]` or omitted | Thinker only | Text |
 | `["audio"]` | Thinker + Talker | Audio (speech) |
 | `["text", "audio"]` | Thinker + Talker | Text + Audio |
+| `["image"]` | Thinker + Imagegen (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
 
 For ready-to-copy curl examples (text / audio / multimodal input, SSE
 streaming, reasoning mode), see the recipe at

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -7,14 +7,15 @@ PORT="${PORT:-8091}"
 QUERY_TYPE="${1:-text}"
 
 # Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then
+if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
     echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]"
+    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
     echo "  text: Text-only query"
     echo "  use_audio: Audio + Text query"
     echo "  use_image: Image + Text query"
     echo "  use_video: Video + Text query"
     echo "  use_mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
     exit 1
 fi
 
@@ -100,6 +101,43 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
+if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
+    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
+    # The shipped ming_flash_omni_dual.yaml runs the thinker on cards 0-3
+    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
+    # section in README.md for knobs and img2img.
+    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
+    request_body=$(cat <<'EOF'
+{
+  "model": "Jonathan1909/Ming-flash-omni-2.0",
+  "modalities": ["image"],
+  "messages": [
+    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
+  ]
+}
+EOF
+)
+    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+        -X POST "http://localhost:${PORT}/v1/chat/completions" \
+        -H "Content-Type: application/json" \
+        -d "$request_body")
+    # Server returns base64-encoded PNG either as a data URL in
+    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
+    b64=$(echo "$output" | jq -r '.choices[0].message.content
+        | if type == "array"
+            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
+            elif type == "string" then . else "" end')
+    if [[ -z "$b64" ]]; then
+        echo "Error: no image returned. Raw response:"
+        echo "$output" | jq '.'
+        exit 1
+    fi
+    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
+    printf '%s' "$b64" | base64 --decode > "$out_path"
+    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
+    exit 0
+fi
+
 request_body=$(cat <<EOF
 {
   "model": "Jonathan1909/Ming-flash-omni-2.0",

pyproject.toml

@@ -249,6 +249,7 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
     ".*[Nn][Dd]3.*",
+    "byt5.*",
 ]
 extend-ignore-re = [
     "\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -261,3 +262,4 @@ Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
 ND = "ND"
+byt5 = "byt5"

vllm_omni/diffusion/data.py

@@ -891,6 +891,17 @@ class OmniDiffusionConfig:
                     self.model_class_name = "SenseNovaU1Pipeline"
                     self.tf_model_config = TransformerConfig()
                     self.update_multimodal_support()
+                elif "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
+                    "bailingmm_moe_v2_lite",
+                    "ming_flash_omni",
+                    "ming_flash_omni_thinker",
+                ):
+                    # Ming-flash-omni-2.0 — imagegen stage uses the custom
+                    # ``MingImagePipeline`` (ZImage DiT + Qwen2 connector). See
+                    # vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py.
+                    self.model_class_name = "MingImagePipeline"
+                    self.tf_model_config = TransformerConfig()
+                    self.update_multimodal_support()
                 elif model_type == "nextstep":
                     if self.model_class_name is None:
                         self.model_class_name = "NextStep11Pipeline"

vllm_omni/diffusion/forward_context.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import vllm.ir
 from vllm.config import VllmConfig
@@ -11,6 +12,9 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
+if TYPE_CHECKING:
+    import torch
+
 
 @dataclass
 class ForwardContext:
@@ -23,6 +27,8 @@ class ForwardContext:
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
+    # Per-request reference latent for img2img DiT models (e.g. Ming)
+    ref_latent: 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
@@ -171,3 +177,14 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:
         _forward_context.denoise_step_idx = step_idx
+
+
+def set_forward_context_ref_latent(ref_latent: torch.Tensor | None) -> None:
+    """Set the per-request reference latent on the active ForwardContext.
+
+    Used by img2img-capable DiT models (e.g. Ming-flash-omni-2.0) so the
+    transformer can read the reference latent from request scope instead of
+    module instance state.
+    """
+    if _forward_context is not None:
+        _forward_context.ref_latent = ref_latent

vllm_omni/diffusion/models/ming_flash_omni/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+from .byte5_encoder import MingByT5Encoder
+from .condition_encoder import MingConditionEncoder
+from .ming_zimage_transformer import MingZImageTransformer2DModel
+from .pipeline_ming_imagegen import (
+    MingImagePipeline,
+    get_ming_image_post_process_func,
+)
+from .t5_block_mapper import T5EncoderBlockByT5Mapper
+
+__all__ = [
+    "MingByT5Encoder",
+    "MingConditionEncoder",
+    "MingImagePipeline",
+    "MingZImageTransformer2DModel",
+    "T5EncoderBlockByT5Mapper",
+    "get_ming_image_post_process_func",
+]

vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py

@@ -0,0 +1,223 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""ByT5 glyph/text encoder for Ming-flash-omni-2.0 image generation.
+
+Ported from Ming's ``load_byte5_and_byte5_tokenizer`` +
+``T5EncoderBlockByT5Mapper`` + ``get_condition_embeds_for_image_gen`` byte5
+branch. The released checkpoint's byte5 weights were trained with font/color
+special tokens, so we replicate that vocabulary extension before loading —
+otherwise ``byte5_model.pt`` would shape-mismatch at the embedding layer.
+
+Typical forward: takes a list of user-supplied prompt strings (possibly with
+``<cn-font-N>`` / ``<color-N>`` markers), returns ``[B, byte5_max_length,
+diffusion_c_input_dim]`` features ready to be concatenated onto cap_feats
+along the sequence dimension.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, T5ForConditionalGeneration
+
+from vllm_omni.diffusion.models.ming_flash_omni.t5_block_mapper import (
+    T5EncoderBlockByT5Mapper,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _add_multilingual_special_tokens(
+    tokenizer,
+    text_encoder: nn.Module,
+    font_ann_path: Path,
+    color_ann_path: Path,
+    add_font: bool,
+    add_color: bool,
+    add_align: bool = False,
+) -> None:
+    """Extend the byte5 vocab with per-language font + color markers.
+
+    Mirrors ``add_special_token_multilingual`` in Ming's bizgen utils. The token
+    set must match the one the checkpoint was trained with, otherwise the
+    resized embedding table shape won't line up with the shipped weights.
+    """
+    idx_font_dict = json.loads(font_ann_path.read_text())
+    idx_color_dict = json.loads(color_ann_path.read_text())
+
+    font_tokens: list[str] = []
+    for font_code in idx_font_dict:
+        prefix = font_code[:3]
+        if prefix in ("cn-", "en-", "jp-", "kr-"):
+            font_tokens.append(f"<{prefix}font-{idx_font_dict[font_code]}>")
+        else:
+            font_tokens.append(f"<font-{idx_font_dict[font_code]}>")
+    color_tokens = [f"<color-{i}>" for i in range(len(idx_color_dict))]
+    align_tokens = [f"<align-{i}>" for i in range(3)]
+
+    extra: list[str] = []
+    if add_color:
+        extra += color_tokens
+    if add_font:
+        extra += font_tokens
+    if add_align:
+        extra += align_tokens
+    tokenizer.add_tokens(extra, special_tokens=True)
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+
+class MingByT5Encoder(nn.Module):
+    """Bundles byte5 tokenizer + T5 encoder + T5EncoderBlockByT5Mapper.
+
+    Build with ``MingByT5Encoder.from_checkpoint(<model>/byte5)`` when the
+    checkpoint ships byte5 weights; otherwise callers can skip this and the
+    pipeline falls back to no-byte5 conditioning.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        text_encoder: nn.Module,
+        mapper: T5EncoderBlockByT5Mapper,
+        max_length: int,
+    ) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.mapper = mapper
+        self.max_length = max_length
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        byte5_dir: Path,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MingByT5Encoder:
+        # Wrap the whole load in fork_rng so any nn.init.normal_ inside
+        # T5ForConditionalGeneration.from_pretrained / vocab-resize cannot
+        # advance the default CPU/CUDA generator state. Without this, the
+        # diffusion pipeline's seeded noise becomes order-dependent across
+        # requests (same-seed replays diverge).
+        cuda_devs = list(range(torch.accelerator.device_count())) if torch.cuda.is_available() else []
+        with torch.random.fork_rng(devices=cuda_devs, enabled=True):
+            return cls._from_checkpoint_impl(byte5_dir, device=device, dtype=dtype)
+
+    @classmethod
+    def _from_checkpoint_impl(
+        cls,
+        byte5_dir: Path,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> MingByT5Encoder:
+        byte5_dir = Path(byte5_dir)
+        # Ming checkpoint uses ``byt5`` (no 'e') in filenames and JSON keys
+        # variable names below keep the ``byte5_`` spelling for readability.
+        cfg_path = byte5_dir / "byt5.json"
+        cfg_raw = json.loads(cfg_path.read_text())
+        cfg = SimpleNamespace(**cfg_raw)
+        byte5_config = cfg.byt5_config
+        mapper_config = cfg.byt5_mapper_config
+        max_length = int(cfg.byt5_max_length)
+
+        # ---- Tokenizer + T5 encoder (base).
+        ckpt_key = byte5_config.get("byt5_ckpt_path")
+        byte5_ckpt_path = byte5_dir / ckpt_key.lstrip("./")
+        tokenizer = AutoTokenizer.from_pretrained(byte5_ckpt_path, local_files_only=True)
+        text_encoder = T5ForConditionalGeneration.from_pretrained(byte5_ckpt_path, local_files_only=True).get_encoder()
+
+        # ---- Extend vocab with font/color markers so the shipped weights load.
+        if byte5_config.get("special_token"):
+            if not byte5_config.get("multilingual", True):
+                raise NotImplementedError(
+                    "Non-multilingual byte5 vocab extension is not ported; "
+                    "the released Ming checkpoint uses multilingual=True."
+                )
+            _add_multilingual_special_tokens(
+                tokenizer,
+                text_encoder,
+                font_ann_path=byte5_dir / byte5_config["font_ann_path"].lstrip("./"),
+                color_ann_path=byte5_dir / byte5_config["color_ann_path"].lstrip("./"),
+                add_font=bool(byte5_config.get("font_special_token")),
+                add_color=bool(byte5_config.get("color_special_token")),
+            )
+
+        # ---- Load byt5 text-encoder weights (base.pt wraps the backbone in a
+        # trainable-module container; byt5_model.pt has the top-level encoder
+        # state). Follow Ming's two-step load.
+        base_state = torch.load(byte5_dir / "byt5_model" / "base.pt", map_location="cpu", weights_only=False)
+        prefix = "module.text_tower.encoder."
+        base_filtered = {
+            name[len(prefix) :]: state for name, state in base_state["state_dict"].items() if name.startswith(prefix)
+        }
+        text_encoder.load_state_dict(base_filtered, strict=True)
+        del base_state, base_filtered
+
+        encoder_state = torch.load(byte5_dir / "byt5_model" / "byt5_model.pt", map_location="cpu", weights_only=False)
+        text_encoder.load_state_dict(encoder_state)
+        del encoder_state
+
+        text_encoder.to(device=device, dtype=dtype).eval()
+
+        # ---- Mapper.
+        mapper = T5EncoderBlockByT5Mapper(
+            byte5_config=text_encoder.config,
+            num_layers=int(mapper_config["num_layers"]),
+            sdxl_channels=int(mapper_config["sdxl_channels"]),
+        )
+        mapper_state = torch.load(byte5_dir / "byt5_mapper" / "byt5_mapper.pt", map_location="cpu", weights_only=False)
+        # mapper now uses vllm-omni TP-aware T5 layers (fused qkv_proj / wi);
+        mapper.load_weights(mapper_state.items())
+        del mapper_state
+        mapper.to(device=device, dtype=dtype).eval()
+
+        logger.info(
+            "[MingByT5Encoder] ready: d_model=%d mapper_layers=%d sdxl_channels=%d max_length=%d vocab=%d",
+            text_encoder.config.d_model,
+            mapper_config["num_layers"],
+            mapper_config["sdxl_channels"],
+            max_length,
+            len(tokenizer),
+        )
+        return cls(tokenizer, text_encoder, mapper, max_length)
+
+    @torch.inference_mode()
+    def forward(self, texts: list[str]) -> torch.Tensor:
+        """Tokenize → T5 encode → mapper; masks out padded positions.
+
+        Returns ``[B, max_length, sdxl_channels]``. Padded positions are zeroed
+        so the downstream ``torch.cat`` with cap_feats doesn't inject garbage.
+        """
+        device = next(self.text_encoder.parameters()).device
+        dtype = next(self.text_encoder.parameters()).dtype
+
+        tokens = self.tokenizer(
+            texts,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = tokens.input_ids.to(device)
+        attention_mask = tokens.attention_mask.to(device)
+
+        encoder_out = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask.float(),
+        )
+        hidden_states = encoder_out[0]
+        feats = self.mapper(hidden_states, attention_mask)
+        feats = feats * attention_mask.unsqueeze(-1).to(dtype=feats.dtype)
+        return feats.to(dtype=dtype)
+
+
+__all__ = ["MingByT5Encoder"]

vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+#
+# Adapted from Ming repository (inclusionAI/Ming) — the ``get_condition_embeds
+# _for_image_gen`` path in modeling_bailingmm2.py.
+
+"""Ming-flash-omni-2.0 condition encoder for image generation.
+
+Pipeline (runs inside the imagegen stage):
+
+    thinker hidden states at query-token positions     [B, N, 4096]
+                         │
+                         ▼ proj_in (Linear, bias=True)
+                                                       [B, N, 1536]
+                         │
+                         ▼ Qwen2ForCausalLM connector (is_causal=False)
+                           loaded from <checkpoint>/connector/
+                                                       [B, N, 1536]
+                         │
+                         ▼ proj_out (Linear, bias=True)
+                                                       [B, N, 2560]
+                         │
+                         ▼ F.normalize(dim=-1) × 1000  (text_encoder_norm)
+                                                       [B, N, 2560]
+                         │
+                         ▼
+            cap_feats consumed by ZImageTransformer2DModel
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MingConditionEncoder(nn.Module):
+    """Wraps a Qwen2 connector + norm/projection, producing DiT condition embeds.
+
+    The connector is a ``Qwen2ForCausalLM`` loaded from the ``connector/``
+    subfolder of the Ming checkpoint. We run its base model in a non-causal
+    (bidirectional) mode, because the connector is used as an encoder over the
+    pre-baked query-token hidden states, not as an autoregressive decoder.
+
+    Args:
+        image_gen_config: ``MingImageGenConfig`` from ``MingFlashOmniConfig``.
+        thinker_hidden_size: Hidden size of the thinker (BailingMoeV2) model.
+            Used to build a ``proj_in`` layer when the connector embedding
+            dim differs. For the released checkpoint this is 4096.
+        device: Placement for the module.
+        dtype: Parameter dtype (typically bfloat16 / float16).
+    """
+
+    def __init__(
+        self,
+        image_gen_config: MingImageGenConfig,
+        *,
+        thinker_hidden_size: int = 4096,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> 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
+
+        # Populated lazily by ``load_from_checkpoint`` to keep this module
+        # cheap to construct (useful for dummy-init paths and unit tests).
+        self.connector: nn.Module | None = None
+        self.connector_hidden_size: int | None = None
+        self.proj_in: nn.Module = nn.Identity()
+        self.proj_out: nn.Module = nn.Identity()
+        self.norm: nn.Module = nn.Identity()
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_from_checkpoint(self, model_path: str | Path) -> None:
+        """Load the Qwen2 connector + optional projection/norm weights.
+
+        This uses HF transformers directly (not vllm's weight loader) because
+        the connector is small (~1.5B params) and only runs once per request
+        as an encoder — vllm's distributed loading machinery is overkill.
+        """
+        from transformers import AutoConfig, Qwen2ForCausalLM
+
+        model_path = Path(model_path)
+        connector_path = model_path / self.config.connector_subfolder
+        logger.info("[MingConditionEncoder] loading connector from %s", connector_path)
+
+        connector_cfg = AutoConfig.from_pretrained(connector_path, trust_remote_code=True, local_files_only=True)
+        # Disable causal masking: the connector is used as a bidirectional
+        # encoder over query-token hidden states in image-gen mode.
+        connector_cfg.is_decoder = False
+        self.connector_hidden_size = int(connector_cfg.hidden_size)
+        logger.info(
+            "[MingConditionEncoder] connector hidden_size=%d, layers=%d",
+            self.connector_hidden_size,
+            getattr(connector_cfg, "num_hidden_layers", -1),
+        )
+
+        connector = Qwen2ForCausalLM.from_pretrained(
+            connector_path,
+            config=connector_cfg,
+            torch_dtype=self._target_dtype,
+            local_files_only=True,
+        )
+        # Force bidirectional attention in every self-attn module. Qwen2 layers
+        # do not branch on ``is_causal`` directly, so we monkey-patch the flag
+        # on each attention block (defensive — some transformers versions read
+        # ``self_attn.is_causal`` in forward).
+        patched = 0
+        for module in connector.modules():
+            if hasattr(module, "is_causal"):
+                module.is_causal = False
+                patched += 1
+        logger.info("[MingConditionEncoder] disabled is_causal on %d sub-modules", patched)
+
+        # We only need the base encoder (no LM head).
+        base = getattr(connector, "model", connector)
+        self.connector = base
+
+        # proj_in: align thinker_hidden_size -> connector_hidden_size.
+        # Observed layout of mlp/ on the released checkpoint
+        # (inclusionAI/Ming-flash-omni-2.0):
+        #   proj_in.{weight,bias}   : Linear(4096 -> 1536)
+        #   proj_out.{weight,bias}  : Linear(1536 -> 2560)
+        #   query_tokens_dict.16x16 : learnable tokens consumed on the thinker
+        #                             side (NOT loaded here).
+        logger.info(
+            "[MingConditionEncoder] adding proj_in: %d -> %d (bias=True)",
+            self.thinker_hidden_size,
+            self.connector_hidden_size,
+        )
+        self.proj_in = nn.Linear(self.thinker_hidden_size, self.connector_hidden_size, bias=True)
+
+        # NOTE: ``text_encoder_norm=True`` in Ming's mlp/config.json refers to
+        # applying ``F.normalize(..., dim=-1)`` (L2 normalization) on the
+        # FINAL cap_feats AFTER proj_out — NOT an intermediate RMSNorm. See
+        # modeling_bailingmm2.py::get_condition_embeds_for_image_gen. We keep
+        # ``self.norm = nn.Identity`` and apply the L2 normalize explicitly
+        # at the end of ``forward``.
+        self.norm = nn.Identity()
+
+        # proj_out: align connector_hidden_size -> diffusion_c_input_dim.
+        c_out = self.config.diffusion_c_input_dim
+        logger.info(
+            "[MingConditionEncoder] adding proj_out: %d -> %d (bias=True)",
+            self.connector_hidden_size,
+            c_out,
+        )
+        self.proj_out = nn.Linear(self.connector_hidden_size, c_out, bias=True)
+
+        mlp_path = model_path / self.config.mlp_subfolder
+        mlp_cfg_path = mlp_path / "config.json"
+        if mlp_cfg_path.exists() and not json.loads(mlp_cfg_path.read_text()).get("use_identity_mlp", False):
+            # MingImagePipeline skips ZImageModel_withMLP, only correct when its inner MLP is Identity.
+            raise NotImplementedError(f"{mlp_cfg_path} has use_identity_mlp=False; ToClipMLP path not implemented.")
+        self._load_optional_mlp_weights(mlp_path)
+
+        if self._target_device is not None:
+            self.to(self._target_device)
+        if self._target_dtype is not None:
+            self.to(dtype=self._target_dtype)
+
+        logger.info(
+            "[MingConditionEncoder] ready: in=%d -> conn=%d -> out=%d",
+            self.thinker_hidden_size,
+            self.connector_hidden_size,
+            c_out,
+        )
+
+    def _load_optional_mlp_weights(self, mlp_path: Path) -> None:
+        """Load proj_in / proj_out (and optional norm) weights from mlp/.
+
+        Expected keys — observed on inclusionAI/Ming-flash-omni-2.0:
+            proj_in.weight           [1536, 4096]
+            proj_in.bias             [1536]
+            proj_out.weight          [2560, 1536]
+            proj_out.bias            [2560]
+            query_tokens_dict.16x16  [256, 4096]   -> thinker-side, skipped here
+            (optional) norm.weight   [1536]
+
+        Any extra keys are logged as warnings so we don't silently miss new
+        fields. Missing proj weights are fatal: they are required for the
+        condition path to be meaningful.
+        """
+        if not mlp_path.exists():
+            logger.warning(
+                "[MingConditionEncoder] mlp/ subfolder missing at %s — proj/norm "
+                "will stay randomly initialized. EXPECT BAD IMAGES until this "
+                "is fixed on real hardware.",
+                mlp_path,
+            )
+            return
+
+        try:
+            from safetensors.torch import load_file  # type: ignore
+        except ImportError:
+            logger.exception("[MingConditionEncoder] safetensors not installed")
+            return
+
+        candidates = sorted(mlp_path.glob("*.safetensors"))
+        if not candidates:
+            candidates = sorted(mlp_path.glob("*.bin"))
+        if not candidates:
+            logger.warning("[MingConditionEncoder] no weight files under %s", mlp_path)
+            return
+
+        state: dict[str, torch.Tensor] = {}
+        for p in candidates:
+            logger.info("[MingConditionEncoder] reading mlp weights: %s", p)
+            if p.suffix == ".safetensors":
+                state.update(load_file(str(p)))
+            else:
+                state.update(torch.load(str(p), map_location="cpu"))
+        logger.info(
+            "[MingConditionEncoder] mlp/ keys: %s",
+            sorted(state.keys()),
+        )
+
+        handled: set[str] = set()
+
+        def _copy(dst: nn.Parameter | torch.Tensor, src_key: str) -> bool:
+            src = state.get(src_key)
+            if src is None:
+                logger.error("[MingConditionEncoder] mlp/ missing key %r", src_key)
+                return False
+            if tuple(src.shape) != tuple(dst.shape):
+                logger.error(
+                    "[MingConditionEncoder] mlp/%s shape mismatch: checkpoint=%s, module=%s",
+                    src_key,
+                    tuple(src.shape),
+                    tuple(dst.shape),
+                )
+                return False
+            with torch.no_grad():
+                dst.copy_(src.to(dtype=dst.dtype, device=dst.device))
+            handled.add(src_key)
+            logger.info(
+                "[MingConditionEncoder] loaded mlp/%s -> %s",
+                src_key,
+                tuple(dst.shape),
+            )
+            return True
+
+        ok_w = _copy(self.proj_in.weight, "proj_in.weight")
+        ok_b = _copy(self.proj_in.bias, "proj_in.bias")
+        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):
+            logger.error(
+                "[MingConditionEncoder] proj_in/proj_out NOT fully loaded; diffusion conditioning will be garbage."
+            )
+
+        # Optional norm weight (Ming uses plain RMSNorm; may or may not ship).
+        if "norm.weight" in state and hasattr(self.norm, "weight"):
+            _copy(self.norm.weight, "norm.weight")
+
+        # query_tokens_dict lives in mlp/ but is consumed on the thinker side
+        # — we neither need nor load it here. Just log the shape so the
+        # thinker-side patch knows what to expect.
+        for k, v in state.items():
+            if k.startswith("query_tokens_dict"):
+                logger.info(
+                    "[MingConditionEncoder] thinker-side parameter %s shape=%s (NOT loaded here; thinker must own it)",
+                    k,
+                    tuple(v.shape),
+                )
+                handled.add(k)
+
+        leftover = set(state.keys()) - handled
+        if leftover:
+            logger.warning("[MingConditionEncoder] mlp/ unhandled keys: %s", sorted(leftover))
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    def forward(
+        self,
+        thinker_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Encode thinker hidden states into DiT condition embeddings.
+
+        Args:
+            thinker_hidden_states: ``[B, N, thinker_hidden_size]`` — sliced at
+                the learnable query-token positions by the stage input
+                processor before being passed here.
+            attention_mask: Optional ``[B, N]`` mask. Defaults to all-ones.
+
+        Returns:
+            ``[B, N, diffusion_c_input_dim]`` condition tensor ready for the
+            ZImage transformer's ``cap_feats`` input.
+        """
+        if self.connector is None:
+            raise RuntimeError("MingConditionEncoder.load_from_checkpoint() must be called before forward().")
+        if thinker_hidden_states.dim() != 3:
+            raise ValueError(f"expected [B, N, H], got shape {tuple(thinker_hidden_states.shape)}")
+
+        b, n, _ = thinker_hidden_states.shape
+        logger.debug(
+            "[MingCondEnc] input shape=%s dtype=%s", tuple(thinker_hidden_states.shape), thinker_hidden_states.dtype
+        )
+
+        x = self.proj_in(thinker_hidden_states)  # [B, N, conn_hidden]
+
+        # Ming's ``get_condition_embeds_for_image_gen`` passes a 4D full-ones
+        # attention mask of shape [B, 1, N, N] to the Qwen2 connector,
+        # forcing full bidirectional self-attention over all query positions.
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (b, 1, n, n),
+                dtype=x.dtype,
+                device=x.device,
+            )
+        elif attention_mask.dim() == 2:
+            mask_2d = attention_mask.to(x.dtype)
+            attention_mask = mask_2d[:, None, None, :].expand(b, 1, n, n)
+
+        out = self.connector(
+            inputs_embeds=x,
+            attention_mask=attention_mask,
+            use_cache=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        hidden = out.hidden_states[-1]  # [B, N, conn_hidden]
+
+        cap_feats = self.proj_out(hidden)  # [B, N, diffusion_c_input_dim]
+
+        # 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.config.text_encoder_norm:
+            cap_feats = cap_feats * 1000.0
+
+        logger.debug("[MingCondEnc] output shape=%s dtype=%s", tuple(cap_feats.shape), cap_feats.dtype)
+        return cap_feats
+
+    # ------------------------------------------------------------------
+    # Utility
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def zero_negative(
+        self,
+        cap_feats: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return a zero tensor shaped like ``cap_feats`` for CFG negatives."""
+        return torch.zeros_like(cap_feats)
+
+    def extra_repr(self) -> str:
+        return (
+            f"thinker_hidden_size={self.thinker_hidden_size}, "
+            f"connector_hidden_size={self.connector_hidden_size}, "
+            f"diffusion_c_input_dim={self.config.diffusion_c_input_dim}, "
+            f"text_encoder_norm={self.config.text_encoder_norm}"
+        )
+
+
+__all__ = ["MingConditionEncoder"]

vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""Ming-specific subclass of ZImageTransformer2DModel that supports ``ref_x``.
+
+Ming's img2img path concatenates a VAE-encoded reference latent along the
+frame axis before patchification, then drops the reference portion from the
+unpatchified output. This is a surgical override — everything else (attention,
+RoPE, final layer) stays on the parent implementation.
+
+"""
+
+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 MingZImageTransformer2DModel(ZImageTransformer2DModel):
+    """ZImage DiT with Ming's reference-latent support."""
+
+    def forward(
+        self,
+        x: list[torch.Tensor],
+        t,
+        cap_feats: list[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+    ):
+        ref_latent = get_forward_context().ref_latent if is_forward_context_available() else None
+        if ref_latent is not None:
+            per_item = ref_latent[0].unsqueeze(1).to(dtype=x[0].dtype, device=x[0].device)  # [C, 1, H, W]
+            x = [torch.cat([img, per_item], dim=1) for img in x]
+        return super().forward(x, t, cap_feats, patch_size=patch_size, f_patch_size=f_patch_size)
+
+    def unpatchify(
+        self,
+        x: list[torch.Tensor],
+        size: list[tuple],
+        patch_size,
+        f_patch_size,
+    ) -> list[torch.Tensor]:
+        out = super().unpatchify(x, size, patch_size, f_patch_size)
+        # No-op when F==1 (pure t2i); drops the reference-frame prediction when F==2.
+        return [t[:, :1, :, :] for t in out]
+
+
+__all__ = ["MingZImageTransformer2DModel"]

vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py

@@ -0,0 +1,520 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""Ming-flash-omni-2.0 imagegen (text-to-image / img2img) diffusion pipeline.
+
+Cross-stage data flow:
+
+    Stage 0 (thinker, llm)           Stage 1 (imagegen, diffusion)
+    ────────────────────             ──────────────────────────────
+    forward returns                  thinker2imagegen hook slices
+    multimodal_output[               final_hidden_states at
+      "final_hidden_states"]         <imagePatch> positions,
+       ↓                             returns list[dict] with
+    shared_memory_connector          {"extra": {"thinker_hidden_states"}}
+       ↓                             ──── via OmniMsgpackEncoder ────>
+                                     MingImagePipeline.forward(req):
+                                       hidden = req.prompts[0]["extra"][...]
+                                       cond = condition_encoder(hidden)
+                                       img = ZImagePipeline-style loop
+                                       return DiffusionOutput(output=img)
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import set_forward_context_ref_latent
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.models.ming_flash_omni.byte5_encoder import (
+    MingByT5Encoder,
+)
+from vllm_omni.diffusion.models.ming_flash_omni.condition_encoder import (
+    MingConditionEncoder,
+)
+from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import (
+    MingZImageTransformer2DModel,
+)
+from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _ZPipelineSamplingParams:
+    """Typed shim satisfying the attributes ZImagePipeline.forward reads
+    from ``req.sampling_params``.  Unlike SimpleNamespace this fails loudly
+    at construction if a required field is missing."""
+
+    height: int
+    width: int
+    num_inference_steps: int
+    guidance_scale: float
+    generator: torch.Generator | None = None
+    strength: float | None = None
+    sigmas: list[float] | None = None
+    max_sequence_length: int = 512
+    guidance_rescale: float | None = None
+    num_outputs_per_prompt: int = 1
+
+
+@dataclass
+class _ZPipelineRequest:
+    """Typed shim satisfying the attributes ZImagePipeline.forward reads
+    from ``req``."""
+
+    request_id: str
+    sampling_params: _ZPipelineSamplingParams
+    prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}])
+
+
+class MingImagePipeline(ZImagePipeline):
+    """Ming-flash-omni-2.0 text-to-image diffusion pipeline.
+
+    Ming-specific components added on top of the inherited contract:
+      * ``condition_encoder`` — Qwen2 connector + proj_in/out + F.normalize×1000
+      * ``byte5``             — Optional ByT5 glyph encoder (loaded if checkpoint
+                                ships ``byt5/``)
+    """
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",  # noqa: ARG002
+    ) -> None:
+        # Skip ZImagePipeline.__init__ (it would eagerly load the Z-Image text
+        # encoder/tokenizer that Ming replaces with its own condition_encoder).
+        nn.Module.__init__(self)
+
+        model_path = od_config.model
+        if not os.path.exists(model_path):
+            model_path = download_weights_from_hf_specific(model_path, od_config.revision, ["*"])
+
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        local_files_only = os.path.exists(model_path)
+
+        self.od_config = od_config
+        self._execution_device = get_local_device()
+        self.device = self._execution_device  # Ming convention alias
+        self._dtype = dtype
+
+        # Ming's per-checkpoint image-gen configuration. We cannot rely on
+        # ``od_config.hf_config.image_gen_config`` because the diffusion
+        # stage is started with ``hf_config_name: thinker_config`` (the
+        # BailingMM2Config), which does not carry a MingImageGenConfig.
+        # Fall back to defaults that match the released checkpoint.
+        self.image_gen_config = MingImageGenConfig()
+        logger.info(
+            "[MingImagePipeline] init: model=%s dtype=%s image_gen_config=%s",
+            model_path,
+            dtype,
+            self.image_gen_config,
+        )
+
+        # ----- weights_sources: DiT transformer + VAE.
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder=self.image_gen_config.transformer_subfolder,
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder=self.image_gen_config.vae_subfolder,
+                revision=od_config.revision,
+                prefix="vae.",
+            ),
+        ]
+
+        prefetch_subfolders(
+            model_path,
+            [self.image_gen_config.scheduler_subfolder, self.image_gen_config.vae_subfolder],
+            local_files_only=local_files_only,
+        )
+
+        # ----- Scheduler: load config-only from disk + Ming-specific override.
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_path,
+            subfolder=self.image_gen_config.scheduler_subfolder,
+            local_files_only=local_files_only,
+        )
+        # Ming forces use_dynamic_shifting=True at runtime regardless of what
+        # the checkpoint scheduler_config.json ships.
+        self.scheduler.config["use_dynamic_shifting"] = True
+        logger.info(
+            "[MingImagePipeline] scheduler: %s (use_dynamic_shifting=True)",
+            type(self.scheduler).__name__,
+        )
+
+        # ----- VAE: DistributedAutoencoderKL.
+        vae_config = DistributedAutoencoderKL.load_config(
+            model_path, subfolder=self.image_gen_config.vae_subfolder, local_files_only=local_files_only
+        )
+        self.vae = DistributedAutoencoderKL.from_config(vae_config).to(self._execution_device, dtype=dtype)
+        self.vae.eval()
+
+        # ----- DiT transformer.
+        self.transformer = MingZImageTransformer2DModel(quant_config=None)
+
+        # Ming brings its own conditioning path — no Z-Image text_encoder /
+        # tokenizer.
+        self.text_encoder = None
+        self.tokenizer = None
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True)
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=getattr(od_config, "enable_diffusion_pipeline_profiler", False)
+        )
+
+        # ----- Condition encoder (Qwen2 connector + proj_in/out + norm×1000).
+        self.condition_encoder = MingConditionEncoder(
+            self.image_gen_config,
+            thinker_hidden_size=self.image_gen_config.thinker_hidden_size,
+            device=self.device,
+            dtype=dtype,
+        )
+        self.condition_encoder.load_from_checkpoint(model_path)
+
+        # Optional ByT5 glyph/text encoder. Only loaded when the checkpoint
+        # ``extra_args.image_gen.byte5_text`` requests will be ignored.
+        byte5_dir = Path(model_path) / "byt5"
+        if byte5_dir.exists():
+            self.byte5 = MingByT5Encoder.from_checkpoint(byte5_dir, device=self.device, dtype=dtype)
+        else:
+            self.byte5 = None
+            logger.info("[MingImagePipeline] no byt5/ subfolder at %s; ByT5 enhancement disabled", byte5_dir)
+
+        logger.info("[MingImagePipeline] ready — vae_scale_factor=%d", self.vae_scale_factor)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _resolve_byte5_texts(extra: dict, sampling_params) -> list[str]:
+        """Resolve byte5 glyph texts.
+
+        Two sources, in order of priority:
+          1. ``extra["byte5_text"]`` — auto-extracted from the user prompt's
+             quoted spans by ``thinker2imagegen`` (already wrapped as
+             ``'Text "<glyph>". '`` by Ming's ``get_text_from_prompt``).
+          2. ``sampling_params.extra_args["image_gen"]["byte5_text"]`` — an
+             explicit override for programmatic callers; raw strings without
+             the ``Text "..."`` wrapper are auto-wrapped here to match the
+             distribution ByT5 was trained on.
+        """
+        # Source 1: auto-extracted, already wrapped. Return as-is if non-empty.
+        raw = extra.get("byte5_text")
+        if isinstance(raw, str):
+            raw = [raw]
+        if isinstance(raw, list):
+            cleaned = [t for t in raw if isinstance(t, str) and t.strip()]
+            if cleaned:
+                return cleaned
+
+        # Source 2: explicit override — wrap raw strings so the byte5 encoder
+        # sees the same ``Text "<glyph>". `` format Ming used during training.
+        raw = ((getattr(sampling_params, "extra_args", None) or {}).get("image_gen") or {}).get("byte5_text")
+        if isinstance(raw, str):
+            raw = [raw]
+        if isinstance(raw, list):
+            out: list[str] = []
+            for t in raw:
+                if not isinstance(t, str):
+                    continue
+                s = t.strip()
+                if not s:
+                    continue
+                # Don't double-wrap if the caller already supplied ``Text "...". ``.
+                out.append(s if s.startswith('Text "') else f'Text "{s}". ')
+            if out:
+                return out
+        return []
+
+    @torch.inference_mode()
+    def _encode_reference_image(self, ref, height: int, width: int) -> torch.Tensor | None:
+        """Turn a PIL/tensor reference image into a VAE latent for ``ref_x``.
+
+        Applies the same shift/scale Ming uses (``(z - shift_factor) * scaling_factor``)
+        so the concatenated frame lives in the DiT's latent space.
+        """
+        if ref is None:
+            return None
+        if not isinstance(ref, torch.Tensor):
+            ref = self.image_processor.preprocess(ref, height, width)
+        ref = ref.to(device=self.device, dtype=self.vae.dtype)
+        latent = self.vae.encode(ref).latent_dist.mode()
+        return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        """Run one text-to-image generation request.
+
+        Args:
+            req: Diffusion request. The cross-stage thinker hidden states
+                must be present at
+                ``req.prompts[0]["extra"]["thinker_hidden_states"]`` as a
+                ``[N, H]`` (or ``[1, N, H]``) tensor, placed there by
+                ``thinker2imagegen``.
+
+        Returns:
+            DiffusionOutput with ``.output`` set to a ``[B, 3, H, W]``
+            image tensor in ``[-1, 1]``. The vllm-omni diffusion engine's
+            output adapter converts this to PIL/base64 downstream.
+        """
+        first_prompt = req.prompts[0] if req.prompts else None
+        if isinstance(first_prompt, str):
+            prompt_dict: dict[str, Any] = {}
+        elif isinstance(first_prompt, dict):
+            prompt_dict = first_prompt
+        elif first_prompt is not None and hasattr(first_prompt, "_asdict"):
+            prompt_dict = first_prompt._asdict()
+        elif first_prompt is not None and hasattr(first_prompt, "__dict__"):
+            prompt_dict = vars(first_prompt)
+        else:
+            prompt_dict = {}
+
+        extra = prompt_dict.get("extra") or {}
+        hidden = extra.get("thinker_hidden_states")
+        if hidden is None:
+            # Same dual-path convention as glm_image: also check
+            # ``sampling_params.extra_args``.
+            hidden = (req.sampling_params.extra_args or {}).get("thinker_hidden_states")
+        if hidden is None:
+            scale = self.image_gen_config.img_gen_scales[-1]
+            num_query_tokens = scale * scale
+            hidden = torch.zeros(
+                (num_query_tokens, self.image_gen_config.thinker_hidden_size),
+                dtype=self._dtype,
+                device=self.device,
+            )
+            logger.warning(
+                "[MingImagePipeline.forward] 'thinker_hidden_states' missing "
+                "from request; falling back to zero-conditioning %s. This is "
+                "expected during warmup; for real requests verify that "
+                "`custom_process_input_func: thinker2imagegen` is set on the "
+                "diffusion stage in the YAML.",
+                tuple(hidden.shape),
+            )
+
+        if not isinstance(hidden, torch.Tensor):
+            raise TypeError(
+                f"[MingImagePipeline] 'thinker_hidden_states' must be a Tensor, got {type(hidden).__name__}"
+            )
+
+        # Move to the pipeline's device+dtype.
+        target_device = next(self.parameters()).device
+        target_dtype = next(self.parameters()).dtype
+        hidden = hidden.to(device=target_device, dtype=target_dtype)
+        if hidden.dim() == 2:
+            hidden = hidden.unsqueeze(0)  # [N, H] -> [1, N, H]
+        logger.debug(
+            "[MingImagePipeline.forward] thinker_hidden_states=%s on %s (%s)",
+            tuple(hidden.shape),
+            target_device,
+            target_dtype,
+        )
+
+        # ----- Condition encoder → cap_feats
+        cap_feats = self.condition_encoder(hidden)
+        logger.debug("[MingImagePipeline.forward] cap_feats=%s", tuple(cap_feats.shape))
+
+        # Real negative CFG conditioning (opt-in). See expand_cfg_prompts.
+        negative_hidden = extra.get("negative_thinker_hidden_states")
+        negative_cap_feats = None
+        if isinstance(negative_hidden, torch.Tensor):
+            negative_hidden = negative_hidden.to(device=target_device, dtype=target_dtype)
+            if negative_hidden.dim() == 2:
+                negative_hidden = negative_hidden.unsqueeze(0)
+            negative_cap_feats = self.condition_encoder(negative_hidden)
+            logger.debug("[MingImagePipeline.forward] negative_cap_feats=%s", tuple(negative_cap_feats.shape))
+
+        # ByT5 text enhancement (opt-in). Appends glyph-aware features along
+        # the sequence dim; negative side gets zeros for the byte5 portion so
+        # CFG doesn't push away from the rendered text.
+        byte5_texts = self._resolve_byte5_texts(extra, req.sampling_params)
+        if byte5_texts and self.byte5 is not None:
+            byte5_feats = self.byte5(byte5_texts).to(device=target_device, dtype=target_dtype)
+            cap_feats = torch.cat((cap_feats, byte5_feats), dim=1)
+            if negative_cap_feats is not None:
+                negative_cap_feats = torch.cat((negative_cap_feats, torch.zeros_like(byte5_feats)), dim=1)
+            logger.debug("[MingImagePipeline.forward] byte5 cat'd: cap_feats=%s", tuple(cap_feats.shape))
+
+        # Sampling knobs: extra_args.image_gen.* > sampling_params.* > MingImageGenConfig defaults.
+        sp = req.sampling_params
+        cfg = self.image_gen_config
+        ig = (sp.extra_args or {}).get("image_gen") or {}
+        resolved: dict[str, Any] = {}
+        for ig_key, sp_attr, default in (
+            ("height", "height", cfg.default_height),
+            ("width", "width", cfg.default_width),
+            ("steps", "num_inference_steps", cfg.num_inference_steps),
+            ("cfg", "guidance_scale", cfg.guidance_scale),
+            ("seed", "seed", None),
+        ):
+            for v in (ig.get(ig_key), getattr(sp, sp_attr), default):
+                if v is not None:
+                    resolved[ig_key] = v
+                    break
+
+        height = int(resolved["height"])
+        width = int(resolved["width"])
+        num_inference_steps = int(resolved["steps"])
+        guidance_scale = float(resolved["cfg"])
+        seed = resolved.get("seed")
+
+        # Always rebuild the generator from the resolved seed. Reusing
+        # ``sp.generator`` causes two problems:
+        #   (1) if the caller pre-seeded it with sp.seed (e.g. the top-level
+        #       ``seed`` key on OmniDiffusionSamplingParams), any override via
+        #       ``extra_args.image_gen.seed`` would be silently ignored; and
+        #   (2) a persistent generator instance accumulates state across
+        #       requests → same-seed replays produce different outputs.
+        if seed is not None:
+            generator = torch.Generator(device=target_device).manual_seed(int(seed))
+        else:
+            generator = sp.generator
+
+        # Format prompt_embeds / negative_prompt_embeds as list[Tensor]
+        # (one entry per request) — matches ZImagePipeline's contract when
+        # prompt_embeds are pre-computed.
+        prompt_embeds = [cap_feats[i] for i in range(cap_feats.shape[0])]
+        if negative_cap_feats is not None:
+            negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])]
+        else:
+            negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds]
+
+        z_sp = _ZPipelineSamplingParams(
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        z_req = _ZPipelineRequest(
+            request_id=req.request_id or "ming-imagegen",
+            sampling_params=z_sp,
+        )
+
+        # Reference image (img2img) → VAE-encoded latent published on the
+        # active ForwardContext so MingZImageTransformer2DModel can read it
+        # from request scope inside its forward().
+        ref_latent = self._encode_reference_image(extra.get("reference_image"), height, width)
+        set_forward_context_ref_latent(ref_latent)
+
+        logger.debug(
+            "[MingImagePipeline.forward] running z_pipeline hw=(%d,%d) steps=%d cfg=%.2f seed=%s overrides=%s ref=%s",
+            height,
+            width,
+            num_inference_steps,
+            guidance_scale,
+            seed,
+            ig,
+            None if ref_latent is None else tuple(ref_latent.shape),
+        )
+        try:
+            output = super().forward(
+                z_req,
+                prompt=None,
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                generator=generator,
+                output_type="pt",
+                return_dict=True,
+            )
+        finally:
+            set_forward_context_ref_latent(None)
+
+        if hasattr(output, "output") and output.output is not None:
+            raw = output.output
+        elif hasattr(output, "images"):
+            raw = output.images
+        else:
+            raw = output
+        if not isinstance(raw, torch.Tensor):
+            raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                "[MingImagePipeline.forward] produced image tensor shape=%s range=[%.3f,%.3f]",
+                tuple(raw.shape),
+                raw.float().min().item(),
+                raw.float().max().item(),
+            )
+        return DiffusionOutput(output=raw)
+
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+def get_ming_image_post_process_func(od_config: OmniDiffusionConfig):
+    """Return a post-process callable that converts the raw VAE tensor to PIL.
+
+    The diffusion engine calls ``post_process_func(output_data)`` where
+    ``output_data`` is the ``DiffusionOutput.output`` tensor returned by
+    ``MingImagePipeline.forward``. It has shape ``[B, 3, H, W]`` in ``[-1, 1]``
+    (Z-image VAE convention). We run the standard ``VaeImageProcessor``
+    postprocess to convert it to ``list[PIL.Image]`` which vllm-omni's
+    ``OmniRequestOutput.from_diffusion`` then bubbles up as
+    ``omni_outputs.images`` for serving_chat to base64-encode.
+
+    Registered via ``_DIFFUSION_POST_PROCESS_FUNCS["MingImagePipeline"]``
+    in vllm_omni/diffusion/registry.py.
+    """
+    import json
+
+    model_path = od_config.model
+    vae_config_path = os.path.join(model_path, "vae", "config.json")
+    try:
+        with open(vae_config_path) as f:
+            vae_cfg = json.load(f)
+        block_out_channels = vae_cfg.get("block_out_channels", [128, 256, 512, 512])
+        vae_scale_factor = 2 ** (len(block_out_channels) - 1)
+    except Exception:
+        vae_scale_factor = 8  # Ming's Flux-format VAE default
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True)
+
+    def post_process_func(images: torch.Tensor):
+        # VaeImageProcessor.postprocess with default output_type="pil"
+        # returns ``list[PIL.Image]``.
+        return image_processor.postprocess(images.float())
+
+    return post_process_func
+
+
+__all__ = [
+    "MingImagePipeline",
+    "get_ming_image_post_process_func",
+]

vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The vLLM-Omni team.
+
+"""T5EncoderBlockByT5Mapper — Ming's per-block T5 stack mapping byte5 features
+onto the DiT condition space.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from diffusers.models.modeling_utils import ModelMixin
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5Block
+
+
+class T5EncoderBlockByT5Mapper(ModelMixin):
+    """Stacks ``num_layers`` T5 encoder blocks on top of byte5 features and
+    projects them to ``sdxl_channels`` (= Ming's ``diffusion_c_input_dim``).
+    """
+
+    def __init__(self, byte5_config, num_layers: int, sdxl_channels: int | None = None) -> None:
+        super().__init__()
+        if num_layers > 0:
+            self.blocks = nn.ModuleList(
+                [
+                    T5Block(
+                        byte5_config,
+                        has_relative_attention_bias=(i == 0),
+                        prefix=f"blocks.{i}",
+                    )
+                    for i in range(num_layers)
+                ]
+            )
+        else:
+            self.blocks = None
+        self.layer_norm = RMSNorm(byte5_config.d_model, eps=byte5_config.layer_norm_epsilon)
+        if sdxl_channels is not None:
+            self.channel_mapper = nn.Linear(byte5_config.d_model, sdxl_channels)
+            self.final_layer_norm = RMSNorm(sdxl_channels, eps=byte5_config.layer_norm_epsilon)
+        else:
+            self.channel_mapper = None
+            self.final_layer_norm = None
+
+    def get_extended_attention_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        if attention_mask.dim() == 3:
+            extended = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            extended = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(f"Unexpected attention_mask shape {tuple(attention_mask.shape)}")
+        extended = extended.to(dtype=dtype)
+        return (1.0 - extended) * torch.finfo(dtype).min
+
+    def forward(self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        extended_mask = self.get_extended_attention_mask(attention_mask, dtype=self.dtype)
+
+        hidden_states = inputs_embeds
+        position_bias = None
+
+        if self.blocks is not None:
+            for block in self.blocks:
+                hidden_states, position_bias = block(
+                    hidden_states,
+                    mask=extended_mask,
+                    position_bias=position_bias,
+                )
+
+        hidden_states = self.layer_norm(hidden_states)
+        if self.channel_mapper is not None:
+            hidden_states = self.channel_mapper(hidden_states)
+            hidden_states = self.final_layer_norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load Ming's HF-format byte5_mapper checkpoint into the fused
+        TP-aware layers.
+
+        Source format (from ``byte5_mapper.pt``):
+            blocks.{i}.layer.0.SelfAttention.{q,k,v,o}.weight
+            blocks.{i}.layer.1.DenseReluDense.{wi_0,wi_1,wo}.weight
+            blocks.{i}.layer.{0,1}.layer_norm.weight
+            {layer_norm, channel_mapper, final_layer_norm}.{weight,bias}
+
+        Target format (after ``T5Block`` from t5_encoder.py):
+            blocks.{i}.layer.0.SelfAttention.qkv_proj.weight  (fused q+k+v)
+            blocks.{i}.layer.1.DenseReluDense.wi.weight       (fused wi_0+wi_1)
+            (others identical)
+        """
+        stacked_params_mapping = [
+            (".qkv_proj.", ".q.", "q"),
+            (".qkv_proj.", ".k.", "k"),
+            (".qkv_proj.", ".v.", "v"),
+            (".DenseReluDense.wi.", ".DenseReluDense.wi_0.", 0),
+            (".DenseReluDense.wi.", ".DenseReluDense.wi_1.", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            matched = False
+            for target_substr, source_substr, shard_id in stacked_params_mapping:
+                if source_substr not in name:
+                    continue
+                target_name = name.replace(source_substr, target_substr, 1)
+                if target_name in params_dict:
+                    param = params_dict[target_name]
+                    param.weight_loader(param, loaded_weight, shard_id)
+                    loaded_params.add(target_name)
+                    matched = True
+                    break
+
+            if matched:
+                continue
+
+            if name in params_dict:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        return loaded_params
+
+
+__all__ = ["T5EncoderBlockByT5Mapper"]

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -832,5 +832,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         loaded_weights = loader.load_weights(weights)
         # Record components loaded by diffusers submodules to satisfy strict checks.
         loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
-        loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+        # downstream pipelines (e.g. MingImagePipeline) may set ``self.text_encoder = None`` when they
+        # bring their own conditioning path.
+        if self.text_encoder is not None:
+            loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
         return loaded_weights

vllm_omni/diffusion/registry.py

@@ -141,6 +141,11 @@ _DIFFUSION_MODELS = {
         "pipeline_bagel",
         "BagelPipeline",
     ),
+    "MingImagePipeline": (
+        "ming_flash_omni",
+        "pipeline_ming_imagegen",
+        "MingImagePipeline",
+    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -467,6 +472,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
+    "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",

vllm_omni/engine/cfg_companion_tracker.py

@@ -22,17 +22,18 @@ class CfgCompanionTracker:
         self._companion_to_parent: dict[str, str] = {}  # companion -> parent
         self._done: dict[str, set[str]] = {}  # parent -> completed companion ids
         self._pending_parents: dict[str, dict[str, Any]] = {}  # parent -> deferred result
+        self._companion_outputs: dict[str, Any] = {}  # companion_id -> engine output
 
     def is_companion(self, req_id: str) -> bool:
         return req_id in self._companion_ids
 
-    def has_companions(self, parent_id: str) -> bool:
-        return parent_id in self._companion_map
-
     def get_parent_id(self, req_id: str) -> str | None:
         """Return the parent request id for a companion, or None."""
         return self._companion_to_parent.get(req_id)
 
+    def has_companions(self, parent_id: str) -> bool:
+        return parent_id in self._companion_map
+
     def all_companions_done(self, parent_id: str) -> bool:
         role_map = self._companion_map.get(parent_id, {})
         done_set = self._done.get(parent_id, set())
@@ -52,6 +53,20 @@ class CfgCompanionTracker:
         self._companion_ids.add(companion_id)
         self._companion_to_parent[companion_id] = parent_id
 
+    def set_companion_output(self, companion_id: str, output: Any) -> None:
+        """Stash companion engine output for the parent to bundle at forward time."""
+        self._companion_outputs[companion_id] = output
+
+    def pop_companion_outputs(self, parent_id: str) -> list[Any]:
+        """Pop companion outputs (role-registration order) for bundling into set_engine_outputs."""
+        role_map = self._companion_map.get(parent_id, {})
+        outputs = []
+        for cid in role_map.values():
+            out = self._companion_outputs.pop(cid, None)
+            if out is not None:
+                outputs.append(out)
+        return outputs
+
     def on_companion_completed(self, companion_id: str) -> str | None:
         """Mark done. Returns parent_id only if parent is pending and all companions finished."""
         parent_id = self._companion_to_parent.get(companion_id)
@@ -86,6 +101,7 @@ class CfgCompanionTracker:
         for companion_id in companion_ids:
             self._companion_ids.discard(companion_id)
             self._companion_to_parent.pop(companion_id, None)
+            self._companion_outputs.pop(companion_id, None)
         self._done.pop(parent_id, None)
         self._pending_parents.pop(parent_id, None)
         return companion_ids

vllm_omni/engine/orchestrator.py

@@ -741,7 +741,10 @@ class Orchestrator:
         finished = output.finished
         submit_ts = req_state.stage_submit_ts.get(stage_id)
 
+        # CFG companion: stash output so parent can bundle [parent, *companions]
+        # into source_outputs for the bridge (e.g. thinker2imagegen).
         if finished and self._cfg_tracker.is_companion(req_id):
+            self._cfg_tracker.set_companion_output(req_id, output)
             await self._handle_cfg_companion_ready(req_id)
             await self._cleanup_request_ids([req_id])
             return
@@ -933,12 +936,34 @@ class Orchestrator:
         requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
 
         if next_pool.stage_type == "diffusion":
+            companion_outputs = self._cfg_tracker.pop_companion_outputs(req_id)
+            expected = len(self._cfg_tracker.get_companion_request_ids(req_id))
+            if expected > len(companion_outputs):
+                logger.warning(
+                    "[Orchestrator] req=%s: only %d/%d CFG companion outputs arrived; "
+                    "downstream CFG conditioning may degrade",
+                    req_id,
+                    len(companion_outputs),
+                    expected,
+                )
+            diffusion_source_outputs = [output, *companion_outputs]
             if next_client.custom_process_input_func is not None:
                 _t_ar2d = _time.perf_counter()
-                diffusion_prompt = next_client.custom_process_input_func(
-                    source_outputs,
+                _fn = next_client.custom_process_input_func
+                _extra_kwargs: dict[str, Any] = {}
+                # TODO: replace signature probe with explicit kwarg contract.
+                try:
+                    import inspect as _inspect
+
+                    if "sampling_params" in _inspect.signature(_fn).parameters:
+                        _extra_kwargs["sampling_params"] = params
+                except (TypeError, ValueError):
+                    pass
+                diffusion_prompt = _fn(
+                    diffusion_source_outputs,
                     req_state.prompt,
                     requires_multimodal_data,
+                    **_extra_kwargs,
                 )
                 _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
                 req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d

vllm_omni/entrypoints/openai/serving_chat.py

@@ -4,6 +4,7 @@ import json
 import time
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator, Callable
+from dataclasses import fields, is_dataclass
 from datetime import datetime, timedelta, timezone
 from io import BytesIO
 from typing import Any, Final, cast
@@ -410,8 +411,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         engine_prompt_image = None
 
                 # Override the prompts produced by chat-template preprocessing.
+                is_img2img = engine_prompt_image is not None
                 tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
-                if engine_prompt_image:
+                if is_img2img:
                     tprompt["modalities"] = ["img2img"]
                 else:
                     tprompt["modalities"] = ["image"]
@@ -428,6 +430,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     mm_processor_kwargs["target_h"] = height
                 if width is not None:
                     mm_processor_kwargs["target_w"] = width
+                # Pass output modalities so model-specific MM processors can
+                # detect image-generation requests and apply their own prompt
+                # rewrites (e.g. query-token expansion, placeholder injection).
+                mm_processor_kwargs["modalities"] = ["img2img"] if is_img2img else ["image"]
                 tprompt["mm_processor_kwargs"] = mm_processor_kwargs
                 if engine_prompt_image is not None:
                     tprompt["multi_modal_data"] = engine_prompt_image
@@ -442,6 +448,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         for k, v in engine_prompt_image.items()
                     }
 
+                # Forward image_gen extra_args from diffusion-stage sampling
+                # params into the prompt dict so that thinker2imagegen can
+                # access them (e.g. byte5_text for ByT5 glyph rendering).
+                if hasattr(request, "sampling_params_list") and request.sampling_params_list:
+                    for sp_raw in request.sampling_params_list:
+                        ea = (
+                            (sp_raw.get("extra_args") or {})
+                            if isinstance(sp_raw, dict)
+                            else (getattr(sp_raw, "extra_args", None) or {})
+                        )
+                        ig = ea.get("image_gen") or {}
+                        if ig:
+                            tprompt["image_gen_extra_args"] = ig
+                            break
+
                 engine_prompts = [tprompt]
                 # Store height/width for applying to diffusion stage sampling params later
                 _image_gen_height = height
@@ -770,13 +791,31 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return new_messages
 
-    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
-        final_sampling_params_list = []
-        for sampling_params in sampling_params_list:
+    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
+        """Convert request dicts to stage-typed sampling params objects.
+
+        For diffusion stages, build ``OmniDiffusionSamplingParams`` so
+        downstream ``StageDiffusionClient._sampling_params_to_dict`` (which
+        requires a dataclass) works. For LLM stages build ``SamplingParams``.
+        """
+        stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
+        final_sampling_params_list: list[Any] = []
+        for idx, sampling_params in enumerate(sampling_params_list):
+            stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm"
+            target_cls = OmniDiffusionSamplingParams if stage_type == "diffusion" else SamplingParams
             if isinstance(sampling_params, dict):
-                final_sampling_params_list.append(SamplingParams(**sampling_params))
-            elif isinstance(sampling_params, SamplingParams):
+                final_sampling_params_list.append(target_cls(**sampling_params))
+            elif isinstance(sampling_params, target_cls):
                 final_sampling_params_list.append(sampling_params)
+            elif isinstance(sampling_params, SamplingParams | OmniDiffusionSamplingParams):
+                # Cross-typed (e.g. user passed SamplingParams but this is a
+                # diffusion stage) — rebuild via a dict round-trip so we end
+                # up with the correct target class.
+                as_dict = {
+                    f.name: getattr(sampling_params, f.name)
+                    for f in (fields(sampling_params) if is_dataclass(sampling_params) else [])
+                } or sampling_params.__dict__
+                final_sampling_params_list.append(target_cls(**as_dict))
             else:
                 raise ValueError(f"Invalid sampling params: {sampling_params}")
         return final_sampling_params_list

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py

@@ -126,15 +126,21 @@ class MingFlashOmniForConditionalGeneration(
             self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors
 
         elif self.model_stage == "imagegen":
-            # TODO: Implement image generator stage
+            # Image generation is a separate diffusion stage; it does not run
+            # through this AR wrapper. See
+            # ``vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py``
+            # and ``model_arch: MingImagePipeline`` in the stage YAML.
             raise NotImplementedError(
-                "Image generation stage is not yet implemented. Please use model_stage='thinker' for now."
+                "Image generation stage is not implemented in this AR wrapper. "
+                "Use stage_type: diffusion with model_arch: MingImagePipeline "
+                "for the imagegen stage (see stage_configs/ming_flash_omni.yaml)."
             )
 
         else:
             raise ValueError(
                 f"Invalid model_stage: {self.model_stage!r}. Must be one of: 'thinker', 'imagegen'. "
-                f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly."
+                f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly. "
+                f"For image generation use stage_type: diffusion with model_arch: MingImagePipeline."
             )
 
     def forward(
@@ -177,9 +183,43 @@ class MingFlashOmniForConditionalGeneration(
         raise NotImplementedError("get_mrope_input_positions not available on current stage")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stripped = ((name.removeprefix("thinker."), value) for name, value in weights)
-        thinker_loaded = self.thinker.load_weights(stripped)
-        return add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+        loaded_weights: set[str] = set()
+        thinker_weights: list[tuple[str, torch.Tensor]] = []
+
+        for name, value in weights:
+            if name.startswith("thinker."):
+                thinker_weights.append((name, value))
+            elif name.startswith("imagegen."):
+                # Imagegen weights are loaded by MingImagePipeline (diffusion
+                # stage) from its own checkpoint subfolders; silently drop
+                # them here so the AR wrapper does not try to load them.
+                continue
+            elif name.startswith("talker."):
+                # Talker weights are loaded by
+                # MingFlashOmniTalkerForConditionalGeneration directly; drop.
+                continue
+            else:
+                # Weights without prefix go to thinker by default
+                thinker_weights.append((name, value))
+
+        if self.model_stage == "thinker" and thinker_weights:
+            stripped = ((name.removeprefix("thinker."), value) for name, value in thinker_weights)
+            thinker_loaded = self.thinker.load_weights(stripped)
+            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
+            loaded_weights.update(thinker_loaded)
+
+            # ``query_tokens_dict.*`` parameters are pre-loaded inside
+            # ``MingFlashOmniThinker.__init__`` from
+            # ``<model>/mlp/model.safetensors`` (not the main shard index),
+            # so vllm's default loader does not see them via the standard
+            # ``load_weights`` path. Report them as loaded here to satisfy the
+            # post-load completeness check.
+            query_tokens_module = getattr(self.thinker, "query_tokens_dict", None)
+            if query_tokens_module is not None:
+                for scale_name, _param in query_tokens_module.items():
+                    loaded_weights.add(f"thinker.query_tokens_dict.{scale_name}")
+
+        return loaded_weights
 
     def get_mm_mapping(self) -> MultiModelKeys:
         return MultiModelKeys.from_string_field(

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -6,6 +6,7 @@
 
 """Ming-flash-omni-2.0 Thinker stage implementation (multimodal understanding)."""
 
+import os
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from typing import Annotated, Any
 
@@ -113,7 +114,9 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
         return 1
 
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
-        return {"image": None, "video": None, "audio": None}
+        # ``img2img`` is the serving-layer alias for a reference image in the
+        # image-gen path; treated as a regular image by ``_MingMultiModalDataParser``.
+        return {"image": None, "video": None, "audio": None, "img2img": 1}
 
     def get_mm_max_tokens_per_item(
         self,
@@ -137,7 +140,14 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
 
         return mm_max_tokens
 
-    def get_feature_extractor(self, **kwargs: object) -> MingWhisperFeatureExtractor:
+    def get_feature_extractor(self, **kwargs: object):
+        """Return the audio feature extractor from the processor.
+
+        The processor may be loaded via trust_remote_code paths that return
+        a stock transformers ``WhisperFeatureExtractor`` rather than
+        vllm-omni's subclass ``MingWhisperFeatureExtractor``. We accept both
+        as long as the caller-needed attributes (``sampling_rate``) exist.
+        """
         hf_processor = self.get_hf_processor(**kwargs)
         feature_extractor = hf_processor.audio_processor
         assert isinstance(feature_extractor, MingWhisperFeatureExtractor)
@@ -145,13 +155,56 @@ class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo):
 
     def get_data_parser(self):
         feature_extractor = self.get_feature_extractor()
-        return MultiModalDataParser(
+        return _MingMultiModalDataParser(
             target_sr=feature_extractor.sampling_rate,
             target_channels=self.get_target_channels(),
             expected_hidden_size=self._get_expected_hidden_size(),
         )
 
 
+class _MingMultiModalDataParser(MultiModalDataParser):
+    """Rewrite ``img2img`` → ``image`` in-place so the thinker embeds the ref.
+
+    Matches Ming's upstream (``processing_bailingmm2.py::__call__``): when
+    ``images`` is provided the image_processor runs → pixel_values + image_grid_thw,
+    and ``_expand_image_tokens`` expands ``<IMAGE>`` placeholders in the text
+    to ``<imagePatch>*N``. The same PIL then ALSO goes to the diffusion stage
+    as ``image_gen_pixel_values_reference``.
+
+    On vllm-omni's serving path serving_chat emits
+    ``multi_modal_data={"img2img": PIL}`` (Bagel convention). We mutate the
+    caller's mm_data dict here — moving the PIL from the ``img2img`` key into
+    the ``image`` key — so that:
+
+      * vLLM's renderer validator sees consistent keys in
+        ``mm_data`` and the parser output (``mm_data_items["image"]``)
+      * The thinker-side image path (``_call_hf_processor`` reads
+        ``mm_data["images"]``, ``_get_prompt_updates`` replaces ``<IMAGE>``
+        placeholders, ``_get_mm_fields_config`` declares image field layout)
+        all run normally with the ref PIL acting as a regular image input
+      * ``thinker2imagegen`` still finds the ref PIL — but now under
+        ``multi_modal_data["image"]`` rather than ``["img2img"]``.
+
+    A companion prepend in ``_apply_hf_processor_main`` ensures the prompt
+    text carries a matching ``<IMAGE>`` placeholder when needed.
+    """
+
+    def parse_mm_data(self, mm_data):
+        if "img2img" in mm_data:
+            img2img = mm_data.pop("img2img")
+            if not isinstance(img2img, list):
+                img2img = [img2img]
+            existing = mm_data.get("image")
+            if existing is None:
+                existing_list: list[Any] = []
+            elif isinstance(existing, list):
+                existing_list = list(existing)
+            else:
+                existing_list = [existing]
+            mm_data["image"] = existing_list + img2img
+        return super().parse_mm_data(mm_data)
+
+
 class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmniThinkerProcessingInfo]):
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_images = mm_counts.get("image", 0)
@@ -185,6 +238,8 @@ class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmn
 
         audio_length = int(audio_duration * sample_rate)
 
+        num_img2img = mm_counts.get("img2img", 0)
+
         mm_data: MultiModalDataDict = {
             "image": self._get_dummy_images(
                 width=image_width,
@@ -200,6 +255,13 @@ class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmn
             "audio": [(np.random.randn(audio_length).astype(np.float32), sample_rate) for _ in range(num_audios)],
         }
 
+        if num_img2img > 0:
+            mm_data["img2img"] = self._get_dummy_images(
+                width=image_width,
+                height=image_height,
+                num_images=num_img2img,
+            )
+
         return mm_data
 
 
@@ -396,6 +458,61 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
     ) -> bool:
         return False
 
+    def _apply_hf_processor_main(
+        self,
+        prompt: str | list[int],
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+        *,
+        enable_hf_prompt_update: bool,
+    ):
+        """Apply Ming-specific prompt rewrites before delegating to vllm's
+        default text / multimodal processing pipeline.
+
+        This is the single point in the base MM-processor call graph where we
+        still have BOTH the full prompt and the request-level
+        ``hf_processor_mm_kwargs`` together. Downstream, vllm's MM caching
+        design splits the work into
+        :meth:`_apply_hf_processor_text_only` (prompt, empty kwargs) and
+        :meth:`_apply_hf_processor_mm_only` (dummy text, real kwargs), so a
+        per-request flag like ``modalities`` cannot reach the real
+        tokenization from inside :meth:`_call_hf_processor`.
+
+        For Ming-flash-omni-2.0, image-generation requests carry
+        ``mm_processor_kwargs["modalities"]`` set to ``["image"]`` (t2i) or
+        ``["img2img"]`` (image edit). We expand the prompt here with the
+        ``<image><imagePatch>*N</image>`` query-token block so the expanded
+        form flows through both the tokenization and (eventual) MM paths.
+        For img2img, we also prepend ``<IMAGE>`` so the thinker-side prompt
+        replacement can locate the reference-image placeholder.
+        """
+        modalities = hf_processor_mm_kwargs.get("modalities") or []
+        is_image_gen = "image" in modalities or "img2img" in modalities
+        if isinstance(prompt, str) and is_image_gen:
+            from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
+                DEFAULT_NUM_QUERY_TOKENS,
+                maybe_expand_image_gen_prompt,
+            )
+
+            # img2img: prepend ``<IMAGE>`` so the ref-image placeholder survives
+            # MM caching (when the ref image is cache-warm, the thinker sees 0
+            # missing items and cannot inject the placeholder itself).
+            if "img2img" in modalities:
+                prompt = "<IMAGE>" + prompt
+
+            ig = getattr(self.info.ctx.model_config.hf_config, "image_gen_config", None)
+            num_query_tokens = getattr(ig, "num_query_tokens", DEFAULT_NUM_QUERY_TOKENS)
+            prompt = maybe_expand_image_gen_prompt(prompt, num_query_tokens=int(num_query_tokens))
+
+        return super()._apply_hf_processor_main(
+            prompt=prompt,
+            mm_items=mm_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+            enable_hf_prompt_update=enable_hf_prompt_update,
+        )
+
     def _call_hf_processor(
         self,
         prompt: str,
@@ -408,6 +525,11 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
         We call the image/audio sub-processors directly (instead of going
         through `MingFlashOmniProcessor.__call__`) so that the high-level
         placeholder tokens remain **unexpanded** in the tokenized output.
+
+        NOTE: Image-gen prompt expansion is handled in
+        :meth:`_apply_hf_processor_main` — by the time we get here, vllm's
+        MM caching has already split text/mm processing and ``mm_kwargs``
+        is empty on the text-only branch.
         """
         hf_processor = self.info.get_hf_processor()
         tokenizer = self.info.get_tokenizer()
@@ -425,6 +547,8 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
 
         videos = mm_data.get("videos", None)
         if videos is not None:
+            # TODO: ``videos=`` on image_processor is deprecated since
+            # transformers v4.57 (removed in v5); migrate to Qwen2VLVideoProcessor.
             video_outputs = hf_processor.image_processor(
                 images=None,
                 videos=videos,
@@ -487,8 +611,13 @@ class MingFlashOmniThinkerForConditionalGeneration(
             return "<VIDEO>"
         elif modality.startswith("audio"):
             return "<AUDIO>"
+        elif modality == "img2img":
+            # Ming's img2img ref image is consumed only by the diffusion stage
+            # (``thinker2imagegen`` reads the raw PIL). No text-side placeholder
+            # is emitted for the thinker.
+            return None
 
-        raise ValueError("Only image, video, or audio modality is supported")
+        raise ValueError("Only image, video, audio, or img2img modality is supported")
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -548,10 +677,79 @@ class MingFlashOmniThinkerForConditionalGeneration(
             )
         logger.info("Initialized WhisperAudioEncoder and AudioProjector")
 
+        # Image-generation learnable query tokens. Stored as a ParameterDict
+        # keyed by scale name (e.g. "16x16" -> [256, hidden_size]) to mirror
+        # the upstream Ming layout in ``mlp/query_tokens_dict.<scale>``. These
+        # are substituted in at ``<imagePatch>`` token positions during
+        # text-to-image requests; image-comprehension requests continue to
+        # fill those positions with vision features as before.
+        self.query_tokens_dict = nn.ParameterDict()
+        self._load_image_gen_query_tokens(vllm_config.model_config.model)
+
         # Expose interfaces
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
-        logger.info("MingFlashOmniThinker initialized with vision and audio towers")
+        logger.info(
+            f"MingFlashOmniThinker initialized with: "
+            f"vision={'yes' if self.vision else 'no'}, "
+            f"audio={'yes' if self.audio else 'no'}, "
+            f"image_gen_scales={list(self.query_tokens_dict.keys())}"
+        )
+
+    def _load_image_gen_query_tokens(self, model_path: str) -> None:
+        """Load learnable query tokens used for image generation.
+
+        Expected layout (observed on inclusionAI/Ming-flash-omni-2.0):
+            <model_path>/mlp/model.safetensors
+                key ``query_tokens_dict.<scale>`` -> [num_tokens, hidden_size]
+                e.g. ``query_tokens_dict.16x16`` -> [256, 4096]
+
+        Missing files are logged as a warning and image-gen is disabled
+        (text-only inference still works).
+        """
+        from pathlib import Path
+
+        p = Path(model_path) / "mlp" / "model.safetensors"
+        if not p.exists() and not os.path.exists(model_path):
+            from vllm_omni.model_executor.model_loader.weight_utils import (
+                download_weights_from_hf_specific,
+            )
+
+            try:
+                local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"])
+                p = Path(local_path) / "mlp" / "model.safetensors"
+            except Exception:
+                logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path)
+        if not p.exists():
+            logger.warning(
+                "[MingFlashOmniThinker] mlp/model.safetensors not found at %s; "
+                "text-to-image will be disabled (text-only inference still works).",
+                p,
+            )
+            return
+        try:
+            from safetensors.torch import load_file  # type: ignore
+
+            state = load_file(str(p))
+        except Exception:
+            logger.exception("[MingFlashOmniThinker] failed to read %s", p)
+            return
+        loaded = 0
+        for key, tensor in state.items():
+            if not key.startswith("query_tokens_dict."):
+                continue
+            scale_name = key.split(".", 1)[1]
+            param = nn.Parameter(tensor.clone(), requires_grad=False)
+            self.query_tokens_dict[scale_name] = param
+            loaded += 1
+            logger.info(
+                "[MingFlashOmniThinker] loaded query_tokens_dict[%s] shape=%s dtype=%s",
+                scale_name,
+                tuple(tensor.shape),
+                tensor.dtype,
+            )
+        if loaded == 0:
+            logger.warning("[MingFlashOmniThinker] no query_tokens_dict.* keys in mlp/model.safetensors")
 
     def extract_image_feature(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:
         """Extract and project image features.
@@ -732,14 +930,86 @@ class MingFlashOmniThinkerForConditionalGeneration(
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.model.word_embeddings(input_ids)
 
-        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
-            return inputs_embeds
+        # Step 1: if vision/audio features were extracted, merge them into
+        # inputs_embeds at their <patch> positions first. This fills ref-image
+        # and other normal vision-input positions with real features.
+        if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:
+            assert is_multimodal is not None, "`is_multimodal` mask required when `multimodal_embeddings` provided"
+            inputs_embeds = _merge_multimodal_embeddings(
+                inputs_embeds=inputs_embeds,
+                multimodal_embeddings=multimodal_embeddings,
+                is_multimodal=is_multimodal,
+            )
 
-        assert is_multimodal is not None, "`is_multimodal` mask required when `multimodal_embeddings` provided"
-        return _merge_multimodal_embeddings(
-            inputs_embeds=inputs_embeds,
-            multimodal_embeddings=multimodal_embeddings,
-            is_multimodal=is_multimodal,
+        # Step 2: image-gen query tokens are appended at the tail of the prompt
+        # as a contiguous block of ``<imagePatch>`` tokens. Inject the learnable
+        # ``query_tokens_dict`` entries at those positions. Running AFTER step 1
+        # is safe because the query-token block is disjoint from ref-image
+        # patch positions (ref patches are earlier in the prompt, filled by
+        # step 1; query tokens occupy the trailing block and get overwritten
+        # here). Injecting always — including when multimodal embeddings are
+        # present — matches Ming's img2img-edit flow where BOTH a real image
+        # input AND the learnable query tokens need to reach the LLM.
+        if len(self.query_tokens_dict) > 0:
+            self._maybe_inject_image_gen_query_embeds(input_ids, inputs_embeds)
+        return inputs_embeds
+
+    def _maybe_inject_image_gen_query_embeds(
+        self,
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+    ) -> None:
+        """Overwrite the trailing ``<imagePatch>`` block with learnable query tokens.
+
+        The image-gen query-token block is appended at the tail of the prompt
+        by ``maybe_expand_image_gen_prompt`` as literal text
+        ``<image><imagePatch>*N</image>``. After tokenization:
+
+            ... <image_start_token> <imagePatch>*N <image_end_token>
+
+        is the exact trailing pattern. We verify that signature before
+        injecting — without the check, a pure vision-understanding request
+        whose final ref-image block ends in N or more ``<imagePatch>`` tokens
+        would be misidentified as image-gen and its vision features silently
+        overwritten. Decode-phase calls (``input_ids`` is one next-token)
+        naturally fail the signature check and return early.
+        """
+        if input_ids.numel() == 0:
+            return
+        patch_token_id = self.config.image_patch_token
+        image_end_token_id = getattr(self.config, "image_end_token", None)
+        if image_end_token_id is None:
+            return
+
+        scale_items = list(self.query_tokens_dict.items())
+        if not scale_items:
+            return
+        scale_name, query_param = scale_items[-1]
+        num_query_tokens = int(query_param.shape[0])
+
+        flat = input_ids.view(-1)
+        L = int(flat.numel())
+        # Need at least num_query_tokens patch tokens + 1 <image_end> token.
+        if L < num_query_tokens + 1:
+            return
+
+        # Signature: the prompt must end with <image_end_token>, and the
+        # N positions immediately before it must all be <imagePatch>.
+        if int(flat[-1].item()) != image_end_token_id:
+            return
+        tail_start = L - 1 - num_query_tokens
+        tail_end = L - 1  # exclusive — last index is <image_end_token>
+        if not (flat[tail_start:tail_end] == patch_token_id).all():
+            return
+
+        query_embeds = query_param.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+        inputs_embeds[tail_start:tail_end] = query_embeds
+        logger.debug(
+            "[MingFlashOmniThinker] injected %d image-gen query embeddings (scale=%s) at positions [%d, %d)",
+            num_query_tokens,
+            scale_name,
+            tail_start,
+            tail_end,
         )
 
     def forward(

vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py

@@ -3,12 +3,86 @@
 # Copyright (c) Ant Group. All rights reserved.
 # Adapted from Ming repo's usage cookbook:
 # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb
-"""Shared prompt-building helpers for Ming-flash-omni standalone talker."""
+"""Ming-flash-omni-2.0 prompt utilities.
+
+Two unrelated families of helpers live in the same file because they
+both belong to the Ming module and are tightly coupled to Ming-specific
+prompt conventions:
+
+1. **Image-gen query-token expansion** — string-level helpers used by
+   the thinker stage to mark the ``<image><imagePatch>*N</image>`` block
+   that gets substituted with learnable ``query_tokens_dict`` embeddings
+   during forward.
+
+2. **TTS / talker caption builder** — JSON caption template for the
+   Ming-flash-omni standalone talker (TTS) stage.
+"""
+
+from __future__ import annotations
 
 import copy
 import json
 from typing import Any
 
+# ============================================================
+# Image-gen query-token block (thinker stage)
+# ============================================================
+
+# Ming's thinker uses these tokens to mark a learnable image-generation
+# query block inside the text prompt. The thinker substitutes its
+# ``query_tokens_dict`` embeddings at each ``<imagePatch>`` position
+# during forward; see
+# ``MingFlashOmniThinker._maybe_inject_image_gen_query_embeds``.
+_IMAGE_OPEN_TOKEN = "<image>"
+_IMAGE_CLOSE_TOKEN = "</image>"
+IMAGE_PATCH_TOKEN = "<imagePatch>"
+
+# Default query-token count matches ``MingImageGenConfig(img_gen_scales=[16])``
+# (16 * 16 = 256), which is what the released inclusionAI/Ming-flash-omni-2.0
+# checkpoint ships.
+DEFAULT_NUM_QUERY_TOKENS = 256
+
+
+def maybe_expand_image_gen_prompt(
+    prompt: str,
+    num_query_tokens: int = DEFAULT_NUM_QUERY_TOKENS,
+) -> str:
+    """Append the ``<image><imagePatch>*N</image>`` suffix for text-to-image.
+
+    The thinker expects image-generation requests to end with an N-wide
+    block of ``<imagePatch>`` tokens (wrapped in ``<image>``/``</image>``)
+    — those positions get substituted with learnable
+    ``query_tokens_dict`` embeddings during forward.
+
+    This helper is a no-op (returns the input unchanged) when:
+
+      * ``prompt`` is not a non-empty string, or
+      * the prompt already contains an ``<imagePatch>`` block (avoids
+        double expansion for tests / manual calls that pre-format the
+        prompt).
+
+    Args:
+        prompt: Raw user prompt text.
+        num_query_tokens: Total number of query tokens to emit. Defaults
+            to 256 (the released checkpoint's ``img_gen_scales=[16]``).
+
+    Returns:
+        The (possibly expanded) prompt text.
+    """
+    if not isinstance(prompt, str) or not prompt:
+        return prompt
+    if IMAGE_PATCH_TOKEN in prompt:
+        return prompt
+
+    # TODO(multi-scale): single-block emission assumes img_gen_scales=[16].
+    suffix = _IMAGE_OPEN_TOKEN + (IMAGE_PATCH_TOKEN * num_query_tokens) + _IMAGE_CLOSE_TOKEN
+    return prompt + suffix
+
+
+# ============================================================
+# TTS / talker caption builder
+# ============================================================
+
 DEFAULT_PROMPT = "Please generate speech based on the following description.\n"
 
 BASE_CAPTION_TEMPLATE: dict[str, Any] = {
@@ -48,3 +122,13 @@ def create_instruction(user_input: dict[str, Any]) -> str:
         if key in item:
             item[key] = value
     return json.dumps(caption, ensure_ascii=False)
+
+
+__all__ = [
+    "IMAGE_PATCH_TOKEN",
+    "DEFAULT_NUM_QUERY_TOKENS",
+    "maybe_expand_image_gen_prompt",
+    "DEFAULT_PROMPT",
+    "BASE_CAPTION_TEMPLATE",
+    "create_instruction",
+]

vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml

@@ -0,0 +1,98 @@
+# Dual-stage single-node deployment of Ming-flash-omni-2.0 text-to-image.
+#
+# Stage 0 (thinker, llm):  BailingMoeV2 LLM + vision/audio encoders, TP=2
+#                          on cards 0 and 1. Detects text-to-image requests
+#                          (auto-appended with <image><imagePatch>*256</image>
+#                          by serving_chat), replaces those positions with
+#                          learnable ``query_tokens_dict`` embeddings, runs
+#                          forward, and exports the final-layer hidden states
+#                          under ``multimodal_output["final_hidden_states"]``.
+#
+# Stage 1 (imagegen, diffusion): Qwen2 connector + ZImage DiT + VAE, TP=1 on
+#                                card 3. Consumes the thinker hidden states
+#                                shipped from stage 0 (sliced at <imagePatch>
+#                                positions by ``thinker2imagegen``) as DiT
+#                                conditioning, runs the Ming condition
+#                                encoder + 30-step FlowMatchEuler denoise
+#                                + VAE decode, emits a raw image tensor.
+#
+# Cross-stage transport: ``shared_memory_connector`` (intra-node). The
+# thinker's ``multimodal_output`` flows automatically; ``thinker2imagegen``
+# (custom_process_input_func) slices it and produces a list[dict] with the
+# conditioning tensor under ``extra[thinker_hidden_states]``. Msgpack
+# (``OmniMsgpackEncoder``) handles bf16/fp16 tensors natively.
+#
+# TODO(ming-flash-omni-2.0): stage 0 `thinker_config` / stage 1 `image_gen_config` target the
+# future unified MingFlashOmniConfig; today's BailingMM2Config checkpoint silently falls back.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    # CFG prompt expansion: when a request carries ``negative_prompt``, one
+    # companion is spawned so both positive and negative hidden states reach
+    # the diffusion stage (consumed by MingImagePipeline as real CFG negatives).
+    prompt_expand_func: vllm_omni.model_executor.stage_input_processors.ming_flash_omni.expand_cfg_prompts
+    hf_config_name: thinker_config
+    runtime:
+      devices: "0,1,2,3"
+    engine_args:
+      model_stage: thinker
+      max_num_seqs: 2  # parent + 1 CFG-text companion (negative prompt)
+      model_arch: MingFlashOmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.85
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      tensor_parallel_size: 4
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 1  # image-gen request: prefill only, no AR decode
+      seed: 42
+      detokenize: false
+    output_connectors:
+      to_stage_1: shared_memory_connector
+
+  - stage_id: 1
+    stage_type: diffusion
+    hf_config_name: image_gen_config
+    runtime:
+      devices: "4"
+    engine_args:
+      model_stage: dit
+      model_arch: MingImagePipeline  # resolved in vllm_omni/diffusion/registry.py
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.ming_flash_omni.thinker2imagegen
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+    input_connectors:
+      from_stage_0: shared_memory_connector
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    shared_memory_connector:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536  # 64KB threshold
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py

@@ -4,27 +4,458 @@
 
 from __future__ import annotations
 
+import logging
+import re
+from dataclasses import dataclass
 from typing import Any
 
+import torch
 from vllm.inputs import TextPrompt
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
+logger = logging.getLogger(__name__)
+
+
+CFG_TEXT_SUFFIX = "__cfg_text"
+
+
+# Fallback when stage config introspection fails; matches
+# llm_config.image_patch_token on the released Ming-flash-omni-2.0 checkpoint.
+_DEFAULT_IMAGE_PATCH_TOKEN_ID = 157157
+
+
+# Ming's byte5 glyph text is auto-extracted from the user prompt's quoted
+# spans (ASCII double quotes / Chinese curly quotes). Patterns and regex
+# taken verbatim from Ming ``processing_bailingmm2.py::get_text_from_prompt``.
+_GLYPH_QUOTE_PATTERNS = [r"\"(.*?)\"", r"‘(.*?)’", r"“(.*?)”"]
+_GLYPH_REMOVE_KEYWORDS = ("remove", "delete", "erase")
+
+
+def _check_single_quotes(s: str) -> bool:
+    """Decide whether ASCII single quotes in ``s`` act as glyph delimiters.
+
+    Mirrors Ming's heuristic: count Chinese chars as weight 3, others as 1;
+    if any paired-single-quote span crosses weight 20 we assume the quotes are
+    apostrophes (e.g. "don't"), not glyph markers.
+    """
+    if s.count("'") % 2 != 0:
+        return False
+    positions = [i for i, ch in enumerate(s) if ch == "'"]
+    for i in range(0, len(positions), 2):
+        start, end = positions[i], positions[i + 1]
+        inner = s[start + 1 : end]
+        chinese = sum(1 for c in inner if "\u4e00" <= c <= "\u9fff")
+        total = 3 * chinese + (len(inner) - chinese)
+        if total >= 20:
+            return False
+    return True
+
+
+def _extract_byte5_glyph_text(prompt: str) -> str:
+    """Return ``Ming``-style ``'Text "<glyph>". '`` or ``""`` when no glyph."""
+    if not isinstance(prompt, str) or not prompt:
+        return ""
+    if "'" in prompt and _check_single_quotes(prompt):
+        prompt = prompt.replace("'", '"')
+
+    texts: list[str] = []
+    for pattern in _GLYPH_QUOTE_PATTERNS:
+        texts.extend(re.findall(pattern, prompt))
+
+    if len(texts) == 1:
+        # Treat "remove/delete/erase ..." as a glyph-removal intent, not generation.
+        text_start = min(
+            (prompt.find(q) for q in ('"', "‘", "“") if prompt.find(q) >= 0),
+            default=-1,
+        )
+        lower = prompt.lower()
+        for kw in _GLYPH_REMOVE_KEYWORDS:
+            idx = lower.find(kw)
+            if 0 <= idx < text_start:
+                return ""
+
+    if not texts:
+        return ""
+    # Only the last quoted span is used (Ming's choice; keeps the most recent edit target).
+    return f'Text "{texts[-1]}". '
+
+
+# ---------------------------------------------------------------------------
+# CFG prompt expansion (stage 0: prompt_expand_func)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _CfgExpandedPrompt:
+    """Minimal structural object consumed by ``AsyncOmniEngine._enqueue_cfg_companions``."""
+
+    prompt: dict[str, Any]
+    role: str
+    request_id_suffix: str
+
+    def apply_overrides(self, base_params: Any, base_spl: list[Any]) -> tuple[Any, list[Any]]:
+        return base_params, base_spl
+
+
+def expand_cfg_prompts(
+    prompt: dict[str, Any] | str,
+    sampling_params: Any,
+) -> list[_CfgExpandedPrompt]:
+    """Expand a text-to-image request into one CFG-text companion (opt-in).
+
+    Triggers only when a non-empty
+    ``sampling_params.extra_args["image_gen"]["negative_prompt"]`` is set on
+    the stage-0 params; otherwise returns ``[]`` and the pipeline falls back
+    to zero negative (Ming's default behavior).
+    """
+    if not isinstance(prompt, dict):
+        return []
+    if prompt.get("modalities") != ["image"]:
+        return []
+
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    image_gen_args = extra_args.get("image_gen") or {}
+    negative = image_gen_args.get("negative_prompt")
+    if not isinstance(negative, str) or not negative.strip():
+        return []
+
+    neg_prompt_dict: dict[str, Any] = {
+        "prompt": negative,
+        "modalities": prompt.get("modalities"),
+    }
+    mm_kwargs = prompt.get("mm_processor_kwargs")
+    if mm_kwargs:
+        neg_prompt_dict["mm_processor_kwargs"] = dict(mm_kwargs)
+
+    return [_CfgExpandedPrompt(prompt=neg_prompt_dict, role="cfg_text", request_id_suffix=CFG_TEXT_SUFFIX)]
+
+
+# ---------------------------------------------------------------------------
+# Thinker → imagegen bridge (stage 1: custom_process_input_func)
+# ---------------------------------------------------------------------------
+
+
+def _resolve_num_query_tokens(stage: Any) -> int | None:
+    """Return the image-gen ``num_query_tokens`` from the source stage's config.
+
+    Falls back to 256 (``img_gen_scales=[16]``) when the stage config lacks
+    a ``MingImageGenConfig``. Cached on the stage object for O(1) re-reads.
+    """
+    cached = getattr(stage, "_ming_num_query_tokens", None)
+    if isinstance(cached, int):
+        return cached
+    n = 256  # Ming-flash-omni-2.0 default (img_gen_scales=[16])
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        ig = getattr(hf_config, "image_gen_config", None)
+        resolved = getattr(ig, "num_query_tokens", None)
+        if isinstance(resolved, int) and resolved > 0:
+            n = resolved
+    except AttributeError:
+        pass
+    try:
+        stage._ming_num_query_tokens = n
+    except AttributeError:
+        pass
+    return n
+
+
+def _resolve_image_end_token_id(stage: Any) -> int | None:
+    """Return the ``<image_end>`` token id from *stage*'s HF config, cached on first call."""
+    cached = getattr(stage, "_image_end_token_id", None)
+    if isinstance(cached, int):
+        return cached
+    token_id: int | None = None
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        llm_config = getattr(hf_config, "llm_config", None)
+        resolved = getattr(llm_config, "image_end_token", None)
+        if isinstance(resolved, int):
+            token_id = resolved
+    except AttributeError:
+        pass
+    try:
+        stage._image_end_token_id = token_id
+    except AttributeError:
+        pass
+    return token_id
+
+
+def _resolve_image_patch_token_id(stage: Any) -> int:
+    """Return the ``<imagePatch>`` token id from *stage*'s HF config, cached on first call."""
+    cached = getattr(stage, "_image_patch_token_id", None)
+    if isinstance(cached, int):
+        return cached
+
+    token_id = _DEFAULT_IMAGE_PATCH_TOKEN_ID
+    try:
+        hf_config = stage.vllm_config.model_config.hf_config
+        llm_config = getattr(hf_config, "llm_config", None)
+        resolved = getattr(llm_config, "image_patch_token", None)
+        if isinstance(resolved, int):
+            token_id = resolved
+    except AttributeError:
+        pass
+
+    try:
+        stage._image_patch_token_id = token_id
+    except AttributeError:
+        pass
+    return token_id
+
 
 def _validate_stage_inputs(stage_list, engine_input_source):
-    """Validate stage inputs and return the source engine outputs."""
     if not engine_input_source:
         raise ValueError("engine_input_source cannot be empty")
-
     stage_id = engine_input_source[0]
     if stage_id >= len(stage_list):
         raise IndexError(f"Invalid stage_id: {stage_id}")
-
     stage = stage_list[stage_id]
     if stage.engine_outputs is None:
         raise RuntimeError(f"Stage {stage_id} has no outputs yet")
+    return stage, stage.engine_outputs
 
-    return stage.engine_outputs
+
+def _ensure_list(x) -> list[int]:
+    """Convert ConstantList / tensor-like to plain list."""
+    if hasattr(x, "_x"):
+        return list(x._x)
+    if isinstance(x, list):
+        return x
+    if hasattr(x, "tolist"):
+        return x.tolist()
+    return list(x)
+
+
+def _slice_patch_hidden(
+    thinker_output: Any,
+    image_patch_token_id: int,
+    tag: str,
+    num_query_tokens: int | None = None,
+    image_end_token_id: int | None = None,
+) -> torch.Tensor | None:
+    """Return ``[N, H]`` hidden at the image-gen query-token block.
+
+    The image-gen block is always appended at the prompt tail as
+    ``<image_start><imagePatch>*N<image_end>`` by
+    ``maybe_expand_image_gen_prompt``. We find the exact trailing window by
+    anchoring on the final ``image_end_token`` and walking back N positions,
+    verifying they're all ``image_patch_token``. Without that signature check
+    we'd risk slicing ref-image patches (img2img case) or comprehension-only
+    patch blocks.
+
+    Falls back to "all patch positions" if ``image_end_token_id`` or
+    ``num_query_tokens`` is not provided — matches the pre-img2img t2i-only
+    behavior used by older callers.
+    """
+    output = thinker_output.outputs[0]
+    mm_out = getattr(output, "multimodal_output", None) or {}
+    full_hidden = mm_out.get("final_hidden_states")
+    if full_hidden is None:
+        logger.warning("[thinker2imagegen] %s: missing final_hidden_states (keys=%s)", tag, list(mm_out.keys()))
+        return None
+
+    prompt_ids = _ensure_list(thinker_output.prompt_token_ids)
+    prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=full_hidden.device)
+    patch_indices = (prompt_ids_t == image_patch_token_id).nonzero(as_tuple=False).squeeze(-1)
+    total_patches = int(patch_indices.numel())
+    if total_patches == 0:
+        logger.warning("[thinker2imagegen] %s: no <imagePatch> tokens in prompt (len=%d)", tag, len(prompt_ids))
+        return None
+
+    if full_hidden.dim() == 3:
+        assert full_hidden.shape[0] == 1, f"expected batch=1, got {full_hidden.shape}"
+        full_hidden = full_hidden[0]
+    if full_hidden.dim() != 2 or full_hidden.shape[0] != prompt_ids_t.shape[0]:
+        logger.warning(
+            "[thinker2imagegen] %s: hidden shape %s inconsistent with prompt len %d",
+            tag,
+            tuple(full_hidden.shape),
+            prompt_ids_t.shape[0],
+        )
+        return None
+
+    selected_indices = patch_indices
+    if num_query_tokens is not None and image_end_token_id is not None:
+        L = int(prompt_ids_t.numel())
+        if L >= num_query_tokens + 1 and int(prompt_ids_t[-1].item()) == image_end_token_id:
+            tail_start = L - 1 - num_query_tokens
+            tail_end = L - 1  # exclusive
+            tail_slice = prompt_ids_t[tail_start:tail_end]
+            if (tail_slice == image_patch_token_id).all():
+                selected_indices = torch.arange(tail_start, tail_end, dtype=torch.long, device=full_hidden.device)
+            else:
+                logger.warning(
+                    "[thinker2imagegen] %s: tail signature mismatch (expected N patches "
+                    "before <image_end>); falling back to all patch positions",
+                    tag,
+                )
+        else:
+            logger.warning(
+                "[thinker2imagegen] %s: image-gen block signature not found at prompt tail "
+                "(len=%d, last_tok=%s); falling back to all patch positions",
+                tag,
+                L,
+                int(prompt_ids_t[-1].item()) if L else None,
+            )
+
+    hidden = full_hidden[selected_indices].detach().contiguous()
+    if logger.isEnabledFor(logging.DEBUG):
+        f = hidden.float()
+        logger.debug(
+            "[thinker2imagegen] %s sliced=%s (%d of %d patches) mean=%+.4f std=%.4f |x|/tok=%.3f",
+            tag,
+            tuple(hidden.shape),
+            int(selected_indices.numel()),
+            total_patches,
+            f.mean().item(),
+            f.std().item(),
+            f.norm(dim=-1).mean().item(),
+        )
+    return hidden
+
+
+def _resolve_token_ids_from_stage_or_defaults(
+    stage: Any | None,
+) -> tuple[int, int | None, int | None]:
+    """Return (image_patch_token_id, image_end_token_id, num_query_tokens).
+
+    Tries to read from the stage's HF config when available (old API).
+    Falls back to Ming-flash-omni-2.0 defaults when stage is None (new API).
+    """
+    if stage is not None:
+        return (
+            _resolve_image_patch_token_id(stage),
+            _resolve_image_end_token_id(stage),
+            _resolve_num_query_tokens(stage),
+        )
+    # Defaults from Ming-flash-omni-2.0:
+    #   llm_config.image_patch_token = 157157
+    #   llm_config.image_end_token   = 157159
+    #   img_gen_scales=[16] -> 16*16 = 256 query tokens
+    return (_DEFAULT_IMAGE_PATCH_TOKEN_ID, 157159, 256)
+
+
+def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | None:
+    """Read ``byte5_text`` from the diffusion-stage sampling_params.
+
+    Looks up ``sampling_params.extra_args["image_gen"]["byte5_text"]`` (the
+    explicit API surface for ByT5 glyph text). Returns ``None`` if absent or
+    malformed, so callers can fall back to other sources.
+    """
+    if sampling_params is None:
+        return None
+    extra = getattr(sampling_params, "extra_args", None)
+    if not isinstance(extra, dict):
+        return None
+    image_gen = extra.get("image_gen")
+    if not isinstance(image_gen, dict):
+        return None
+    texts = image_gen.get("byte5_text")
+    if isinstance(texts, list) and texts:
+        return [t for t in texts if isinstance(t, str)]
+    return None
+
+
+def thinker2imagegen(
+    source_outputs: list[Any],
+    prompt: Any | None = None,
+    requires_multimodal_data: bool = False,  # noqa: ARG001
+    sampling_params: Any | None = None,
+) -> list[dict[str, Any]]:
+    """Bridge thinker AR outputs into image-generation DiT inputs.
+
+    The orchestrator passes ``source_outputs`` as
+    ``[parent_output, *companion_outputs]``. Parent outputs feed
+    ``extra[thinker_hidden_states]``; the cfg_text companion feeds
+    ``extra[negative_thinker_hidden_states]`` used by MingImagePipeline as real
+    CFG negative conditioning. Unknown-suffix outputs are skipped.
+
+    ``sampling_params`` is the diffusion stage's own SamplingParams, supplied
+    by the orchestrator. ByT5 explicit ``byte5_text`` is read from
+    ``sampling_params.extra_args.image_gen.byte5_text`` (preferred); falls
+    back to ``prompt.image_gen_extra_args.byte5_text``, then to auto-extraction
+    from quoted prompt text.
+    """
+    thinker_outputs = source_outputs
+    image_patch_token_id, image_end_token_id, num_query_tokens = _resolve_token_ids_from_stage_or_defaults(stage=None)
+
+    parent_output = None
+    negative_output = None
+    for o in thinker_outputs:
+        rid = getattr(o, "request_id", "")
+        if rid.endswith(CFG_TEXT_SUFFIX):
+            negative_output = o
+        elif parent_output is None:
+            parent_output = o
+
+    if parent_output is None:
+        logger.warning("[thinker2imagegen] no parent output in engine_outputs; skipping")
+        return []
+
+    parent_hidden = _slice_patch_hidden(
+        parent_output,
+        image_patch_token_id,
+        tag="parent",
+        num_query_tokens=num_query_tokens,
+        image_end_token_id=image_end_token_id,
+    )
+    if parent_hidden is None:
+        return []
+
+    extra: dict[str, Any] = {"thinker_hidden_states": parent_hidden}
+    if negative_output is not None:
+        neg_hidden = _slice_patch_hidden(
+            negative_output,
+            image_patch_token_id,
+            tag="cfg_text",
+            num_query_tokens=num_query_tokens,
+            image_end_token_id=image_end_token_id,
+        )
+        if neg_hidden is not None:
+            extra["negative_thinker_hidden_states"] = neg_hidden
+
+    # img2img: forward the reference image PIL/tensor to the diffusion stage.
+    if isinstance(prompt, dict):
+        mm_data = prompt.get("multi_modal_data") or {}
+        ref_image = mm_data.get("image")
+        if isinstance(ref_image, list) and ref_image:
+            ref_image = ref_image[0]
+        if ref_image is None:
+            ref_image = mm_data.get("img2img")
+            if isinstance(ref_image, list) and ref_image:
+                ref_image = ref_image[0]
+        if ref_image is not None:
+            extra["reference_image"] = ref_image
+
+        # ByT5 glyph text: prefer sampling_params (stage-1 explicit API), then
+        # prompt.image_gen_extra_args (serving_chat path), then auto-extract
+        # from quoted prompt text.
+        prompt_text = prompt.get("prompt", "")
+        sp_byte5 = _extract_byte5_from_sampling_params(sampling_params)
+        if sp_byte5:
+            byte5_texts: list[str] | None = sp_byte5
+        else:
+            ig_extra = prompt.get("image_gen_extra_args") or {}
+            cand = ig_extra.get("byte5_text")
+            byte5_texts = cand if isinstance(cand, list) and cand else None
+
+        if byte5_texts:
+            extra["byte5_text"] = [
+                t if t.startswith("Text ") else f'Text "{t}". ' for t in byte5_texts if isinstance(t, str)
+            ]
+        else:
+            glyph = _extract_byte5_glyph_text(prompt_text)
+            if glyph:
+                extra["byte5_text"] = [glyph]
+    else:
+        # prompt is not a dict — still honor explicit byte5_text from sampling_params.
+        sp_byte5 = _extract_byte5_from_sampling_params(sampling_params)
+        if sp_byte5:
+            extra["byte5_text"] = [t if t.startswith("Text ") else f'Text "{t}". ' for t in sp_byte5]
+
+    return [{"prompt": "", "extra": extra}]
 
 
 def thinker2talker(
@@ -39,7 +470,7 @@ def thinker2talker(
     a talker input prompt with the text and any speaker/instruction info
     from the original request.
     """
-    source_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    _, source_outputs = _validate_stage_inputs(stage_list, engine_input_source)
 
     if not isinstance(prompt, list):
         prompt = [prompt]
@@ -102,3 +533,11 @@ def thinker2talker(
         )
 
     return talker_inputs
+
+
+__all__ = [
+    "CFG_TEXT_SUFFIX",
+    "expand_cfg_prompts",
+    "thinker2imagegen",
+    "thinker2talker",
+]

vllm_omni/transformers_utils/configs/__init__.py

@@ -24,6 +24,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
     "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
+    "MingImageGenConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
@@ -44,6 +45,7 @@ __all__ = [
     "OmniVoiceConfig",
     "BailingMoeV2Config",
     "BailingMM2Config",
+    "MingImageGenConfig",
     "MingFlashOmniConfig",
     "Qwen3VLMoeVisionConfig",
     "WhisperEncoderConfig",

vllm_omni/transformers_utils/configs/ming_flash_omni.py

@@ -260,6 +260,61 @@ class BailingMM2Config(PretrainedConfig):
         return self.llm_config
 
 
+class MingImageGenConfig(PretrainedConfig):
+    """Configuration for Ming-flash-omni-2.0 image generation stage.
+
+    Mirrors the layout of the HF checkpoint at
+    https://huggingface.co/inclusionAI/Ming-flash-omni-2.0 where image-gen
+    components live in sibling subfolders (``connector/``, ``transformer/``,
+    ``vae/``, ``scheduler/``, ``mlp/``).
+    """
+
+    model_type = "ming_flash_omni_imagegen"
+
+    def __init__(
+        self,
+        connector_subfolder: str = "connector",
+        transformer_subfolder: str = "transformer",
+        vae_subfolder: str = "vae",
+        scheduler_subfolder: str = "scheduler",
+        mlp_subfolder: str = "mlp",
+        diffusion_c_input_dim: int = 2560,
+        thinker_hidden_size: int = 4096,
+        img_gen_scales: list[int] | None = None,
+        text_encoder_norm: bool = True,
+        # Default sampling parameters — match Ming's defaults in
+        # ``modeling_bailingmm2.py::generate`` (``image_gen_steps=30``,
+        # ``image_gen_cfg=2.0``).
+        num_inference_steps: int = 30,
+        guidance_scale: float = 2.0,
+        default_height: int = 1024,
+        default_width: int = 1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.connector_subfolder = connector_subfolder
+        self.transformer_subfolder = transformer_subfolder
+        self.vae_subfolder = vae_subfolder
+        self.scheduler_subfolder = scheduler_subfolder
+        self.mlp_subfolder = mlp_subfolder
+        self.diffusion_c_input_dim = diffusion_c_input_dim
+        self.thinker_hidden_size = thinker_hidden_size
+        self.img_gen_scales = img_gen_scales if img_gen_scales is not None else [16]
+        self.text_encoder_norm = text_encoder_norm
+        self.num_inference_steps = num_inference_steps
+        self.guidance_scale = guidance_scale
+        self.default_height = default_height
+        self.default_width = default_width
+
+    @property
+    def num_query_tokens(self) -> int:
+        """Total number of learnable query tokens appended to the thinker input.
+
+        For ``img_gen_scales=[16]`` this yields 256 tokens (a single 16x16 grid).
+        """
+        return sum(s * s for s in self.img_gen_scales)
+
+
 class MingFlashOmniTalkerConfig(PretrainedConfig):
     """Configuration class for Ming-flash-omni-2.0 talker (TTS) stage.
 
@@ -309,13 +364,14 @@ class MingFlashOmniConfig(PretrainedConfig):
     is_composition = True
     sub_configs: ClassVar = {
         "thinker_config": BailingMM2Config,
+        "image_gen_config": MingImageGenConfig,
         "talker_config": MingFlashOmniTalkerConfig,
     }
 
     def __init__(
         self,
         thinker_config: BailingMM2Config | dict[str, Any] | None = None,
-        image_gen_config: dict[str, Any] | None = None,
+        image_gen_config: MingImageGenConfig | dict[str, Any] | None = None,
         talker_config: MingFlashOmniTalkerConfig | dict[str, Any] | None = None,
         **kwargs,
     ):
@@ -326,8 +382,13 @@ class MingFlashOmniConfig(PretrainedConfig):
         else:
             self.thinker_config = thinker_config or BailingMM2Config()
 
-        # Image generation config (for future implementation)
-        self.image_gen_config = image_gen_config
+        # Image generation config
+        if isinstance(image_gen_config, dict):
+            self.image_gen_config = MingImageGenConfig(**image_gen_config)
+        elif image_gen_config is None:
+            self.image_gen_config = MingImageGenConfig()
+        else:
+            self.image_gen_config = image_gen_config
 
         # Talker config
         if isinstance(talker_config, dict):
@@ -342,6 +403,7 @@ class MingFlashOmniConfig(PretrainedConfig):
 # Register model_type -> config class for AutoConfig
 AutoConfig.register(BailingMoeV2Config.model_type, BailingMoeV2Config)
 AutoConfig.register(BailingMM2Config.model_type, BailingMM2Config)
+AutoConfig.register(MingImageGenConfig.model_type, MingImageGenConfig)
 AutoConfig.register(MingFlashOmniTalkerConfig.model_type, MingFlashOmniTalkerConfig)
 AutoConfig.register(MingFlashOmniConfig.model_type, MingFlashOmniConfig)