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