[Config Refactor] Migrate Ming-flash-omni-2.0 Image-Gen deploy configs (#3975)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Co-authored-by: Zheng Wengang <zwg0606@gmail.com>

examples/offline_inference/ming_flash_omni/README.md

@@ -8,7 +8,7 @@ vLLM-Omni supports three deployment modes:
 |------|--------------|--------|
 | 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) |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.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).
 

examples/online_serving/ming_flash_omni/README.md

@@ -10,7 +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 |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.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).
 
@@ -63,34 +63,27 @@ 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:
+Launch with the image deploy 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 \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
     --stage-init-timeout 1800 \
-    --port 8188
+    --init-timeout 1800 \
+    --port 8091
 ```
 
 Then request image output by passing `"modalities": ["image"]`:
 
 ```bash
-curl http://127.0.0.1:8188/v1/chat/completions \
+curl -s http://127.0.0.1:8091/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."}],
+      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
       "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
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
 ```
 
 ### Optional knobs
@@ -109,7 +102,7 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
 Example with all the knobs:
 
 ```bash
-curl http://127.0.0.1:8188/v1/chat/completions \
+curl http://127.0.0.1:8091/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "Jonathan1909/Ming-flash-omni-2.0",
@@ -122,7 +115,7 @@ curl http://127.0.0.1:8188/v1/chat/completions \
             "byte5_text":["理解与生成统一"]}}}
       ],
       "messages": [{"role": "user", "content": "Draw a poster."}]
-    }' -o /tmp/ming_response.json
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
 ```
 
 ### img2img (reference image)
@@ -142,9 +135,9 @@ into the diffusion stage as `extra[reference_image]`:
 
 ### GPU layout
 
-The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
+The shipped `ming_flash_omni_image.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
+`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.
@@ -156,7 +149,7 @@ warmup takes roughly an extra 30–60 s on top of the thinker — set
 | `["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`) |
+| `["image"]` | Thinker + Imagegen (image deploy 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

@@ -103,7 +103,7 @@ 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
+    # The shipped ming_flash_omni_image.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}"

tests/test_config_factory.py

@@ -1580,6 +1580,74 @@ class TestMingFlashOmniPipeline:
         assert len(stages) == 1
         assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
 
+    def test_image_pipeline_registered(self):
+        p = _PIPELINE_REGISTRY.get("ming_flash_omni_image")
+        assert p is not None
+        assert p.model_arch == "MingFlashOmniForConditionalGeneration"
+        assert len(p.stages) == 2
+        assert p.validate() == []
+
+    def test_image_thinker_stage(self):
+        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
+        assert s.model_stage == "thinker"
+        assert s.execution_type == StageExecutionType.LLM_AR
+        assert s.input_sources == ()
+        assert s.final_output is False
+        assert s.owns_tokenizer is True
+        assert s.requires_multimodal_data is True
+        # Image variant exports hidden states for the diffusion stage.
+        assert s.engine_output_type == "latent"
+        assert s.hf_config_name == "thinker_config"
+        assert s.sampling_constraints["detokenize"] is False
+        assert s.prompt_expand_func is not None
+
+    def test_image_dit_stage(self):
+        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
+        assert s.model_stage == "dit"
+        assert s.execution_type == StageExecutionType.DIFFUSION
+        assert s.input_sources == (0,)
+        assert s.final_output is True
+        assert s.final_output_type == "image"
+        assert s.hf_config_name == "image_gen_config"
+        assert s.model_arch == "MingImagePipeline"
+        assert s.custom_process_input_func is not None
+
+    def test_image_processor_wiring_resolves(self):
+        """The prompt_expand_func and custom_process_input_func strings must point to real callables."""
+        thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
+        dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
+        for ref in (thinker.prompt_expand_func, dit.custom_process_input_func):
+            module_path, _, attr = ref.rpartition(".")
+            module = importlib.import_module(module_path)
+            assert callable(getattr(module, attr))
+
+    def test_image_yaml_loads_and_merges(self):
+        """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
+        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
+        if not deploy_path.exists():
+            pytest.skip("ming_flash_omni_image deploy yaml not found")
+
+        deploy = load_deploy_config(deploy_path)
+        assert len(deploy.stages) == 2
+        assert deploy.async_chunk is False
+        assert deploy.pipeline == "ming_flash_omni_image"
+        assert deploy.connectors is not None
+        assert "shared_memory_connector" in deploy.connectors
+
+        pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"]
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        assert len(stages) == 2
+        # Stage 0 thinker: AR worker that emits latents.
+        assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
+        assert stages[0].yaml_engine_args["engine_output_type"] == "latent"
+        assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is False
+        assert stages[0].yaml_extras["prompt_expand_func"] is not None
+        # Stage 1 dit: diffusion stage with MingImagePipeline.
+        assert stages[1].yaml_engine_args["model_arch"] == "MingImagePipeline"
+        assert stages[1].custom_process_input_func is not None
+        assert stages[1].final_output is True
+        assert stages[1].final_output_type == "image"
+
 
 class TestBaseConfigInheritance:
     """Test deploy YAML base_config inheritance."""

vllm_omni/config/pipeline_registry.py

@@ -117,6 +117,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
         "MING_FLASH_OMNI_THINKER_ONLY_PIPELINE",
     ),
+    "ming_flash_omni_image": (
+        "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
+        "MING_FLASH_OMNI_IMAGE_PIPELINE",
+    ),
     "moss_tts_nano": (
         "vllm_omni.model_executor.models.moss_tts_nano.pipeline",
         "MOSS_TTS_NANO_PIPELINE",

vllm_omni/deploy/ming_flash_omni_image.yaml

@@ -0,0 +1,63 @@
+# Ming-flash-omni-2.0 deploy: thinker + image-generation (text-to-image / img2img).
+#
+# Stage 0: Thinker — multimodal understanding -> final-layer hidden states.
+#          TP=4 across GPUs 0-3.
+# Stage 1: DiT (MingImagePipeline) — condition encoder + FlowMatchEuler
+#          denoise + VAE decode. Diffusion worker, runs eager on GPU 4.
+#
+# Cross-stage transport: shared_memory_connector (intra-node).
+#
+# Fields omitted from a stage fall back to StageDeployConfig dataclass
+# defaults (see vllm_omni/config/stage_config.py).
+#   * Stage 0 max_num_seqs=2 leaves room for the CFG-text companion that
+#     `expand_cfg_prompts` (pipeline-level) spawns when a request carries
+#     a negative prompt — positive and negative hidden states both reach
+#     the diffusion stage.
+#   * Image-gen warmup adds latency on top of the thinker;
+#     pass `--stage-init-timeout <seconds>` if the default 300s is tight.
+pipeline: ming_flash_omni_image
+async_chunk: false
+
+trust_remote_code: true
+distributed_executor_backend: mp
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    # max_num_seqs=2: parent + 1 CFG-text companion (negative prompt).
+    max_num_seqs: 2
+    gpu_memory_utilization: 0.85
+    enforce_eager: true
+    tensor_parallel_size: 4
+    devices: "0,1,2,3"
+    output_connectors:
+      to_stage_1: shared_memory_connector
+    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
+
+  - stage_id: 1
+    enforce_eager: true
+    devices: "4"
+    input_connectors:
+      from_stage_0: shared_memory_connector
+    default_sampling_params:
+      seed: 42
+
+connectors:
+  shared_memory_connector:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536  # 64KB threshold
+
+edges:
+  - from: 0
+    to: 1
+    window_size: -1
+    max_inflight: 1
+
+# Ming-flash-omni-2.0 has only been manually verified on CUDA devices
+platforms: {}

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py

@@ -128,12 +128,12 @@ class MingFlashOmniForConditionalGeneration(
         elif self.model_stage == "imagegen":
             # 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.
+            # `vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py`
+            # and `model_arch: MingImagePipeline` in the deploy YAML.
             raise NotImplementedError(
                 "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)."
+                "Use the DIFFUSION execution_type with model_arch: MingImagePipeline "
+                "for the imagegen stage (see vllm_omni/deploy/ming_flash_omni_image.yaml)."
             )
 
         else:

vllm_omni/model_executor/models/ming_flash_omni/pipeline.py

@@ -105,3 +105,37 @@ MING_FLASH_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# Thinker + image-generation (diffusion) variant: text-to-image / img2img.
+MING_FLASH_OMNI_IMAGE_PIPELINE = PipelineConfig(
+    model_type="ming_flash_omni_image",
+    model_arch="MingFlashOmniForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            hf_config_name="thinker_config",
+            engine_output_type="latent",
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            sampling_constraints={"detokenize": False},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            # Resolved in vllm_omni/diffusion/registry.py.
+            model_arch="MingImagePipeline",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            hf_config_name="image_gen_config",
+            custom_process_input_func=f"{_PROC}.thinker2imagegen",
+        ),
+    ),
+)

vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml

@@ -1,98 +0,0 @@
-# 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