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

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

vllm_omni/diffusion/models/ming_flash_omni/byte5_encoder.py

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