From a698269e4712ee31d78c0440aa5a032a93dd94d7 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Wed, 27 May 2026 23:21:20 +0800 Subject: [PATCH] [Model] Add Ming-flash-omni-2.0 Image Generation (Diffusion) Stage (#2875) Signed-off-by: ZhengWG Signed-off-by: Zheng Wengang Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> --- docs/models/supported_models.md | 1 + .../ming_flash_omni/README.md | 14 +- .../online_serving/ming_flash_omni/README.md | 92 ++++ .../run_curl_multimodal_generation.sh | 42 +- pyproject.toml | 2 + vllm_omni/diffusion/data.py | 11 + vllm_omni/diffusion/forward_context.py | 17 + .../models/ming_flash_omni/__init__.py | 20 + .../models/ming_flash_omni/byte5_encoder.py | 223 ++++++++ .../ming_flash_omni/condition_encoder.py | 373 +++++++++++++ .../ming_zimage_transformer.py | 50 ++ .../ming_flash_omni/pipeline_ming_imagegen.py | 520 ++++++++++++++++++ .../models/ming_flash_omni/t5_block_mapper.py | 129 +++++ .../models/z_image/pipeline_z_image.py | 5 +- vllm_omni/diffusion/registry.py | 6 + vllm_omni/engine/cfg_companion_tracker.py | 22 +- vllm_omni/engine/orchestrator.py | 29 +- vllm_omni/entrypoints/openai/serving_chat.py | 51 +- .../models/ming_flash_omni/ming_flash_omni.py | 52 +- .../ming_flash_omni_thinker.py | 294 +++++++++- .../models/ming_flash_omni/prompt_utils.py | 86 ++- .../stage_configs/ming_flash_omni_dual.yaml | 98 ++++ .../stage_input_processors/ming_flash_omni.py | 449 ++++++++++++++- .../transformers_utils/configs/__init__.py | 2 + .../configs/ming_flash_omni.py | 68 ++- 25 files changed, 2614 insertions(+), 42 deletions(-) create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/__init__.py create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py create mode 100644 vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py create mode 100644 vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d26a579ed..03a9a5789 100644 --- a/docs/models/supported_models.md +++ b/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 + imagegen1) | `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md index 4b798619e..83259a021 100644 --- a/examples/offline_inference/ming_flash_omni/README.md +++ b/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). diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md index 533a30e1c..7b538efe8 100644 --- a/examples/online_serving/ming_flash_omni/README.md +++ b/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,"}} + ] +}] +``` + +### 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 diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh index 715f23edd..3b8f4ddc1 100755 --- a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh +++ b/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 <\\.\\.\\.nd", @@ -261,3 +262,4 @@ Vai = "Vai" tockens = "tockens" CANN = "CANN" ND = "ND" +byt5 = "byt5" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a6fe1e4e9..0675c226a 100644 --- a/vllm_omni/diffusion/data.py +++ b/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" diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 35c93c799..f2b6927d2 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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 diff --git a/vllm_omni/diffusion/models/ming_flash_omni/__init__.py b/vllm_omni/diffusion/models/ming_flash_omni/__init__.py new file mode 100644 index 000000000..d10aed623 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py b/vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py new file mode 100644 index 000000000..e0748abcc --- /dev/null +++ b/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 +```` / ```` 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"") + color_tokens = [f"" for i in range(len(idx_color_dict))] + align_tokens = [f"" 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(/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"] diff --git a/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py b/vllm_omni/diffusion/models/ming_flash_omni/condition_encoder.py new file mode 100644 index 000000000..6ebca4f5c --- /dev/null +++ b/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 /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"] diff --git a/vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py b/vllm_omni/diffusion/models/ming_flash_omni/ming_zimage_transformer.py new file mode 100644 index 000000000..4d3925f6d --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py new file mode 100644 index 000000000..0da26aff8 --- /dev/null +++ b/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"] 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 "". '`` 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 "". `` 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", +] diff --git a/vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py b/vllm_omni/diffusion/models/ming_flash_omni/t5_block_mapper.py new file mode 100644 index 000000000..75243529e --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 345ccca65..388b6316c 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index d8302c115..fe7d00c77 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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", diff --git a/vllm_omni/engine/cfg_companion_tracker.py b/vllm_omni/engine/cfg_companion_tracker.py index 1105e887c..aeec9e848 100644 --- a/vllm_omni/engine/cfg_companion_tracker.py +++ b/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 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ae0c82765..48026abab 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index ec10606a1..06ccfcd70 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py index 0111f8069..fd52dcbf8 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py +++ b/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 + # ``/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( diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index bde7477b9..0a9ce0f5c 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/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 ```` placeholders in the text + to ``*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 ```` + 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 ```` 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 + ``*N`` query-token block so the expanded + form flows through both the tokenization and (eventual) MM paths. + For img2img, we also prepend ```` 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 ```` 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 = "" + 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 "