# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from collections.abc import Iterable from enum import Enum from typing import TYPE_CHECKING, Any import torch import torch.nn as nn import torch.nn.functional as F from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig logger = init_logger(__name__) def _positive_divisors(n: int) -> set[int]: if n <= 0: return set() divs: set[int] = set() for d in range(1, int(math.isqrt(n)) + 1): if n % d == 0: divs.add(d) divs.add(n // d) return divs def validate_glm_image_tp_constraints( *, dim: int, num_heads: int, ffn_hidden_dim: int, tensor_parallel_size: int, ) -> list[int]: """Validate GLM-Image TP constraints without requiring a distributed context. Args: dim: Model hidden dimension num_heads: Number of attention heads ffn_hidden_dim: FFN hidden dimension tensor_parallel_size: TP size to validate against Returns: List of supported TP candidates Raises: ValueError: If constraints are violated """ tp_size = int(tensor_parallel_size) if tp_size <= 0: raise ValueError(f"tensor_parallel_size must be > 0, got {tp_size}") if dim % tp_size != 0: supported = sorted(_positive_divisors(dim)) raise ValueError( f"GLM-Image requires dim % tensor_parallel_size == 0, " f"but got dim={dim}, tp={tp_size}. " f"Supported tp candidates by dim: {supported}" ) if num_heads % tp_size != 0: supported = sorted(_positive_divisors(num_heads)) raise ValueError( f"GLM-Image requires num_heads % tensor_parallel_size == 0, " f"but got num_heads={num_heads}, tp={tp_size}. " f"Supported tp candidates by num_heads: {supported}" ) if ffn_hidden_dim % tp_size != 0: supported = sorted(_positive_divisors(ffn_hidden_dim)) raise ValueError( f"GLM-Image requires ffn_hidden_dim % tensor_parallel_size == 0, " f"but got ffn_hidden_dim={ffn_hidden_dim}, tp={tp_size}. " f"Supported tp candidates by ffn_hidden_dim: {supported}" ) supported_tp_candidates = sorted( _positive_divisors(num_heads) & _positive_divisors(dim) & _positive_divisors(ffn_hidden_dim) ) return supported_tp_candidates class GlmImageImageProjector(nn.Module): """Projects latent image patches to transformer hidden dimension.""" def __init__( self, in_channels: int = 16, hidden_size: int = 2560, patch_size: int = 2, ): super().__init__() self.patch_size = patch_size self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, channel, height, width = hidden_states.shape post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p] hidden_states = hidden_states.reshape( batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size ) hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) hidden_states = self.proj(hidden_states) return hidden_states class GlmImageRotaryPosEmbed(nn.Module): """Rotary positional embedding for 2D image patches.""" def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: super().__init__() self.dim = dim self.patch_size = patch_size self.theta = theta def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: batch_size, num_channels, height, width = hidden_states.shape height, width = height // self.patch_size, width // self.patch_size dim_h, dim_w = self.dim // 2, self.dim // 2 h_inv_freq = 1.0 / ( self.theta ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h) ) w_inv_freq = 1.0 / ( self.theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w) ) h_seq = torch.arange(height, device=hidden_states.device) w_seq = torch.arange(width, device=hidden_states.device) h_inv_freq = h_inv_freq.to(hidden_states.device) w_inv_freq = w_inv_freq.to(hidden_states.device) freqs_h = torch.outer(h_seq, h_inv_freq) freqs_w = torch.outer(w_seq, w_inv_freq) # Create position matrices: [height, 1, dim//4] and [1, width, dim//4] freqs_h = freqs_h.unsqueeze(1).expand(height, width, -1) freqs_w = freqs_w.unsqueeze(0).expand(height, width, -1) # Concatenate: [height, width, dim//2] -> [height, width, dim] freqs = torch.cat([freqs_h, freqs_w], dim=-1) freqs = torch.cat([freqs, freqs], dim=-1) freqs = freqs.reshape(height * width, -1) return (freqs.cos(), freqs.sin()) class GlmImagePrepare(nn.Module): """Prepare module for GLM-Image that handles patch embedding and RoPE computation. This module encapsulates the input processing pipeline to create a module boundary where _sp_plan can shard outputs via split_output=True. Similar to Qwen-Image's ImageRopePrepare, this ensures hidden_states and RoPE embeddings are sharded together to maintain dimension alignment. """ def __init__( self, image_projector: nn.Module, rope: GlmImageRotaryPosEmbed, patch_size: int, ): super().__init__() self.image_projector = image_projector self.rope = rope self.patch_size = patch_size def forward( self, hidden_states: torch.Tensor, prior_hidden_states: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Process hidden_states and compute RoPE embeddings. Args: hidden_states: Input latent tensor [B, C, H, W] prior_hidden_states: Optional prior embedding to add Returns: hidden_states: Patched hidden states [B, seq_len, D] rope_cos: RoPE cos embeddings [seq_len, dim] rope_sin: RoPE sin embeddings [seq_len, dim] post_patch_height: Scalar tensor for height after patching post_patch_width: Scalar tensor for width after patching """ batch_size, num_channels, height, width = hidden_states.shape post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) # Compute RoPE (uses original 4D hidden_states shape) image_rotary_emb = self.rope(hidden_states) rope_cos = image_rotary_emb[0].to(hidden_states.device) rope_sin = image_rotary_emb[1].to(hidden_states.device) # Patch embedding: [B, C, H, W] -> [B, seq_len, D] hidden_states = self.image_projector(hidden_states) # Add prior embedding if provided if prior_hidden_states is not None: hidden_states = hidden_states + prior_hidden_states return hidden_states, rope_cos, rope_sin, post_patch_height, post_patch_width class GlmImageAdaLayerNormZero(nn.Module): """Adaptive LayerNorm with zero initialization for both image and text streams.""" def __init__( self, embedding_dim: int, dim: int, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ) -> None: super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.linear = ReplicatedLinear( embedding_dim, 12 * dim, bias=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.linear", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor ) -> tuple[torch.Tensor, ...]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) emb = self.linear(temb) if isinstance(emb, tuple): emb = emb[0] ( shift_msa, c_shift_msa, scale_msa, c_scale_msa, gate_msa, c_gate_msa, shift_mlp, c_shift_mlp, scale_mlp, c_scale_mlp, gate_mlp, c_gate_mlp, ) = emb.chunk(12, dim=1) hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1) return ( hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) class GlmImageAdaLayerNormContinuous(nn.Module): """Final AdaLN for output projection (no activation before Linear).""" def __init__( self, embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine: bool = True, eps: float = 1e-5, bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.linear = ReplicatedLinear( conditioning_embedding_dim, embedding_dim * 2, bias=bias, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.linear", ) self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine) def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: # NO SiLU here emb = self.linear(conditioning_embedding.to(x.dtype)) if isinstance(emb, tuple): emb = emb[0] scale, shift = torch.chunk(emb, 2, dim=1) x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] return x class KVCacheMode(Enum): """Mode for KV cache operations. - WRITE: Store the K/V tensors from condition images - READ: Concatenate cached K/V with current K/V - SKIP: Do not use cache (pass-through) """ WRITE = "write" READ = "read" SKIP = "skip" class GlmImageLayerKVCache: """KV cache for a single attention layer. Stores key and value tensors for image editing. The cache accumulates KV pairs during write mode and provides them during read mode. Shape convention (vllm-omni): key/value: [batch_size, seq_length, num_heads, head_dim] """ def __init__(self): self.k_cache: torch.Tensor | None = None self.v_cache: torch.Tensor | None = None def store(self, key: torch.Tensor, value: torch.Tensor) -> None: """Store or accumulate KV tensors. If cache is empty, stores the tensors directly. If cache is not empty, concatenates new tensors along seq_length dim. Args: key: Key tensor of shape [B, S, H, D] value: Value tensor of shape [B, S, H, D] """ if self.k_cache is None: self.k_cache = key self.v_cache = value else: # Concatenate along sequence dimension (dim=1 for [B, S, H, D]) self.k_cache = torch.cat([self.k_cache, key], dim=1) self.v_cache = torch.cat([self.v_cache, value], dim=1) def get(self) -> tuple[torch.Tensor | None, torch.Tensor | None]: """Get cached KV tensors. Returns: Tuple of (k_cache, v_cache), both may be None if cache is empty. """ return self.k_cache, self.v_cache def clear(self) -> None: """Clear the cache.""" self.k_cache = None self.v_cache = None @property def is_empty(self) -> bool: """Check if cache is empty.""" return self.k_cache is None def __repr__(self) -> str: if self.is_empty: return "GlmImageLayerKVCache(empty)" return f"GlmImageLayerKVCache(k_shape={self.k_cache.shape}, v_shape={self.v_cache.shape})" class GlmImageKVCache: """Container for all layers' KV caches. Manages KV cache for all transformer layers in GLM-Image model. Provides a unified interface for setting mode and clearing cache. Args: num_layers: Number of transformer layers in the model. Example: kv_cache = GlmImageKVCache(num_layers=28) kv_cache.set_mode(KVCacheMode.WRITE) # ... process condition image ... kv_cache.set_mode(KVCacheMode.READ) # ... process target image ... kv_cache.clear() """ def __init__(self, num_layers: int): self.num_layers = num_layers self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] self._mode: KVCacheMode | None = None def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: """Get cache for a specific layer. Args: layer_idx: Index of the layer (0-indexed). Returns: GlmImageLayerKVCache for the specified layer. Raises: IndexError: If layer_idx is out of range. """ if layer_idx < 0 or layer_idx >= self.num_layers: raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})") return self.caches[layer_idx] def __len__(self) -> int: """Return number of layers.""" return self.num_layers @property def mode(self) -> KVCacheMode | None: """Get current cache mode.""" return self._mode def set_mode(self, mode: KVCacheMode | str | None) -> None: """Set cache mode for all layers. Args: mode: Cache mode (WRITE, READ, SKIP) or string ("write", "read", "skip"). Use None to disable cache operations. Raises: ValueError: If mode is an invalid string. """ if mode is None: self._mode = None elif isinstance(mode, str): try: self._mode = KVCacheMode(mode.lower()) except ValueError: raise ValueError(f"Invalid mode: '{mode}', must be one of 'write', 'read', 'skip'") else: self._mode = mode def clear(self) -> None: """Clear cache for all layers and reset mode.""" for cache in self.caches: cache.clear() self._mode = None @property def is_empty(self) -> bool: """Check if all layer caches are empty.""" return all(cache.is_empty for cache in self.caches) def __repr__(self) -> str: mode_str = self._mode.value if self._mode else "None" return f"GlmImageKVCache(num_layers={self.num_layers}, mode={mode_str}, is_empty={self.is_empty})" class GlmImageAttention(nn.Module): """ Joint attention for GLM-Image model using vllm-omni's optimized attention. This combines text and image streams for joint attention computation. Supports KV caching for image editing workflows via external cache. """ def __init__( self, dim: int, num_heads: int, head_dim: int, parallel_config: DiffusionParallelConfig | None = None, out_bias: bool = True, eps: float = 1e-5, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.dim = dim self.total_num_heads = num_heads self.head_dim = head_dim self.parallel_config = parallel_config # QKV projection (fused for efficiency) self.to_qkv = QKVParallelLinear( hidden_size=dim, head_size=head_dim, total_num_heads=num_heads, total_num_kv_heads=num_heads, bias=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_qkv", ) # QK normalization (LayerNorm, not RMSNorm for GLM-Image) self.norm_q = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps) self.norm_k = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps) # Output projection self.to_out = nn.ModuleList( [ RowParallelLinear( dim, dim, bias=out_bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_out.0", ), nn.Dropout(0.0), ] ) # Attention self.attn = Attention( num_heads=self.to_qkv.num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, num_kv_heads=self.to_qkv.num_kv_heads, ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, kv_cache: GlmImageLayerKVCache | None = None, kv_cache_mode: KVCacheMode | None = None, hidden_states_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Forward pass for joint attention. Args: hidden_states: Image hidden states [B, img_seq_len, D] (sharded in SP mode) encoder_hidden_states: Text hidden states [B, text_seq_len, D] (full in SP mode) image_rotary_emb: Tuple of (cos, sin) for RoPE (sharded in SP mode) attention_mask: Optional attention mask kv_cache: Optional layer KV cache for image editing kv_cache_mode: Cache mode (WRITE, READ, SKIP) hidden_states_mask: Mask for SP padding (True=valid, False=padding) Returns: Tuple of (image_hidden_states, text_hidden_states) """ dtype = encoder_hidden_states.dtype batch_size, text_seq_length, _ = encoder_hidden_states.shape # Check if SP is enabled sp_size = self.parallel_config.sequence_parallel_size if self.parallel_config else None use_sp = sp_size is not None and sp_size > 1 if use_sp: forward_ctx = get_forward_context() use_sp = not forward_ctx.split_text_embed_in_sp # Concatenate text and image: [text, image] hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1) # QKV projection qkv = self.to_qkv(hidden_states_combined.contiguous()) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) # Reshape: [B, S, H*D] -> [B, S, H, D] query = query.unflatten(-1, (self.to_qkv.num_heads, -1)) key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) # QK normalization query = self.norm_q(query).to(dtype=dtype) key = self.norm_k(key).to(dtype=dtype) if use_sp: # SP mode: use joint attention mechanism # Split Q/K/V into text and image parts text_query = query[:, :text_seq_length, :, :] text_key = key[:, :text_seq_length, :, :] text_value = value[:, :text_seq_length, :, :] img_query = query[:, text_seq_length:, :, :] img_key = key[:, text_seq_length:, :, :] img_value = value[:, text_seq_length:, :, :] # Apply RoPE only to image part if image_rotary_emb is not None: from diffusers.models.embeddings import apply_rotary_emb img_query = apply_rotary_emb(img_query, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) img_key = apply_rotary_emb(img_key, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) # Create attention metadata for joint attention attn_metadata = AttentionMetadata( joint_query=text_query, joint_key=text_key, joint_value=text_value, joint_strategy="front", ) # Add padding mask for SP if available if hidden_states_mask is not None: attn_metadata.attn_mask = hidden_states_mask # Attention computation with joint text/image # Note: Ulysses post_attention returns [text, image] concatenated joint_hidden_states_out = self.attn(img_query, img_key, img_value, attn_metadata) # Project combined [text, image] outputs, then split. # This keeps SP numerically aligned with the non-SP path. joint_hidden_states_out = joint_hidden_states_out.flatten(2, 3).to(dtype) # Contiguous for FP8/W4A16 quantized RowParallelLinear for module in self.to_out: if isinstance(module, RowParallelLinear): joint_hidden_states_out = module(joint_hidden_states_out.contiguous()) else: joint_hidden_states_out = module(joint_hidden_states_out) encoder_hidden_states_out = joint_hidden_states_out[:, :text_seq_length, :] hidden_states_out = joint_hidden_states_out[:, text_seq_length:, :] else: # Non-SP mode: original logic # Apply RoPE only to image tokens (not text tokens) if image_rotary_emb is not None: query_img = query[:, text_seq_length:, :, :] key_img = key[:, text_seq_length:, :, :] from diffusers.models.embeddings import apply_rotary_emb query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1) key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1) # Handle KV cache for image editing if kv_cache is not None and kv_cache_mode is not None: if kv_cache_mode == KVCacheMode.WRITE: kv_cache.store(key, value) elif kv_cache_mode == KVCacheMode.READ: k_cached, v_cached = kv_cache.get() if k_cached is not None: key = torch.cat([k_cached, key], dim=1) value = torch.cat([v_cached, value], dim=1) # Attention computation attn_metadata = None if attention_mask is not None: if attention_mask.dim() == 3: attention_mask = attention_mask.unsqueeze(1) attn_metadata = AttentionMetadata(attn_mask=attention_mask) hidden_states_out = self.attn(query, key, value, attn_metadata) hidden_states_out = hidden_states_out.flatten(2, 3) hidden_states_out = hidden_states_out.to(dtype) # Output projection for module in self.to_out: if isinstance(module, RowParallelLinear): hidden_states_out = module(hidden_states_out.contiguous()) else: hidden_states_out = module(hidden_states_out) # Split back to text and image encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :] hidden_states_out = hidden_states_out[:, text_seq_length:, :] return hidden_states_out, encoder_hidden_states_out class ColumnParallelGELU(nn.Module): def __init__( self, dim_in: int, dim_out: int, *, approximate: str = "none", bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.proj = ColumnParallelLinear( dim_in, dim_out, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.proj", ) self.approximate = approximate def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) return F.gelu(x, approximate=self.approximate) class ColumnParallelSiLU(nn.Module): def __init__( self, dim_in: int, dim_out: int, *, bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.proj = ColumnParallelLinear( dim_in, dim_out, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.proj", ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) return F.silu(x) class GlmImageFeedForward(nn.Module): def __init__( self, dim: int, dim_out: int | None = None, mult: int = 4, inner_dim: int | None = None, bias: bool = True, activation_fn: str = "gelu", quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() inner_dim = inner_dim or int(dim * mult) dim_out = dim_out or dim if activation_fn == "linear-silu": layers: list[nn.Module] = [ ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0"), nn.Identity(), RowParallelLinear( inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.net.2", ), ] else: approximate = "tanh" if activation_fn == "gelu-approximate" else "none" layers = [ ColumnParallelGELU( dim, inner_dim, approximate=approximate, bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0", ), nn.Identity(), RowParallelLinear( inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.net.2", ), ] self.net = nn.ModuleList(layers) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for module in self.net: if isinstance(module, ColumnParallelLinear): hidden_states, _ = module(hidden_states) elif isinstance(module, RowParallelLinear): hidden_states = module(hidden_states.contiguous()) else: hidden_states = module(hidden_states) return hidden_states class GlmImageTransformerBlock(nn.Module): """Single transformer block for GLM-Image.""" def __init__( self, dim: int = 2560, num_attention_heads: int = 64, attention_head_dim: int = 40, time_embed_dim: int = 512, ffn_hidden_dim: int | None = None, parallel_config: DiffusionParallelConfig | None = None, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ) -> None: super().__init__() # 1. Attention with AdaLN self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim, quant_config=quant_config, prefix=f"{prefix}.norm1") self.attn1 = GlmImageAttention( dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.attn1", ) # 2. Feedforward self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.ff = GlmImageFeedForward( dim=dim, dim_out=dim, inner_dim=ffn_hidden_dim, activation_fn="gelu-approximate", quant_config=quant_config, prefix=f"{prefix}.ff", ) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, attention_kwargs: dict[str, Any] | None = None, kv_cache: GlmImageLayerKVCache | None = None, kv_cache_mode: KVCacheMode | None = None, hidden_states_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Forward pass for transformer block. Args: hidden_states: Image hidden states encoder_hidden_states: Text hidden states temb: Timestep embedding image_rotary_emb: RoPE embeddings attention_mask: Text attention mask attention_kwargs: Additional attention arguments kv_cache: Layer-specific KV cache for image editing kv_cache_mode: Cache mode (WRITE, READ, SKIP) hidden_states_mask: Mask for SP padding (True=valid, False=padding) Returns: Tuple of (image_hidden_states, text_hidden_states) """ # 1. Timestep conditioning via AdaLN ( norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = self.norm1(hidden_states, encoder_hidden_states, temb) # 2. Attention attn_hidden_states, attn_encoder_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states, image_rotary_emb=image_rotary_emb, attention_mask=attention_mask, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, hidden_states_mask=hidden_states_mask, ) hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1) # 3. Feedforward norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * ( 1 + c_scale_mlp.unsqueeze(1) ) + c_shift_mlp.unsqueeze(1) ff_output = self.ff(norm_hidden_states) ff_output_context = self.ff(norm_encoder_hidden_states) hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) return hidden_states, encoder_hidden_states class GlmImageTransformer2DModel(CachedTransformer): """ GLM-Image Transformer model for 2D image generation. This is the vllm-omni optimized version of the GLM-Image DiT model. Args: od_config: OmniDiffusionConfig containing model configuration. Transformer hyper-parameters (e.g. patch size / channels / heads) are read from `od_config.tf_model_config`. """ _repeated_blocks = ["GlmImageTransformerBlock"] # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together. # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where # _sp_plan can shard outputs via split_output=True. # # Key insight: hidden_states and RoPE embeddings MUST be sharded together # to maintain dimension alignment for RoPE computation in attention layers. _sp_plan = { # Shard GlmImagePrepare outputs (hidden_states and RoPE must be sharded together) "prepare": { # hidden_states: [B, seq_len, D] - shard along sequence dimension 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), # RoPE cos: [seq_len, dim] - shard along sequence dimension 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), # RoPE sin: [seq_len, dim] - shard along sequence dimension 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), # post_patch_height and post_patch_width are scalars, not sharded }, # Gather output at proj_out "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } _hsdp_shard_conditions = [is_transformer_block_module] def __init__( self, od_config: OmniDiffusionConfig, quant_config: "QuantizationConfig | None" = None, ): super().__init__() patch_size = od_config.tf_model_config.patch_size in_channels = od_config.tf_model_config.in_channels out_channels = od_config.tf_model_config.out_channels num_attention_heads = od_config.tf_model_config.num_attention_heads attention_head_dim = od_config.tf_model_config.attention_head_dim time_embed_dim = od_config.tf_model_config.time_embed_dim condition_dim = od_config.tf_model_config.condition_dim prior_vq_quantizer_codebook_size = od_config.tf_model_config.prior_vq_quantizer_codebook_size text_embed_dim = od_config.tf_model_config.text_embed_dim # Get num_layers from config if available model_config = od_config.tf_model_config num_layers = getattr(model_config, "num_layers", 28) if model_config is not None else 28 self.od_config = od_config self.patch_size = patch_size self.in_channels = in_channels self.out_channels = out_channels self.parallel_config = od_config.parallel_config # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords pooled_projection_dim = 2 * 2 * condition_dim inner_dim = num_attention_heads * attention_head_dim tp_size = self.parallel_config.tensor_parallel_size ffn_hidden_dim = inner_dim * 4 supported_tp_candidates = validate_glm_image_tp_constraints( dim=inner_dim, num_heads=num_attention_heads, ffn_hidden_dim=ffn_hidden_dim, tensor_parallel_size=tp_size, ) logger.info_once( "GLM-Image init: dim=%d num_heads=%d head_dim=%d ffn_hidden_dim=%d tp=%d (supported_tp=%s)", inner_dim, num_attention_heads, attention_head_dim, ffn_hidden_dim, tp_size, tuple(supported_tp_candidates), ) # 1. RoPE self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0) # 2. Patch & Text-timestep embedding self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size) self.glyph_projector = GlmImageFeedForward( dim=text_embed_dim, dim_out=inner_dim, inner_dim=inner_dim, activation_fn="gelu", quant_config=quant_config, prefix="glyph_projector", ) self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim) self.prior_projector = GlmImageFeedForward( dim=inner_dim, dim_out=inner_dim, inner_dim=inner_dim, activation_fn="linear-silu", quant_config=quant_config, prefix="prior_projector", ) # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan) self.prepare = GlmImagePrepare(self.image_projector, self.rope, patch_size) self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( embedding_dim=time_embed_dim, condition_dim=condition_dim, pooled_projection_dim=pooled_projection_dim, timesteps_dim=time_embed_dim, ) # 3. Transformer blocks self.transformer_blocks = nn.ModuleList( [ GlmImageTransformerBlock( inner_dim, num_attention_heads, attention_head_dim, time_embed_dim, ffn_hidden_dim=ffn_hidden_dim, parallel_config=self.parallel_config, quant_config=quant_config, prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) ] ) # 4. Output projection # Final modulation feeds proj_out; quant_config is NOT applied here # to avoid precision degradation in the final projection layer. self.norm_out = GlmImageAdaLayerNormContinuous( inner_dim, time_embed_dim, elementwise_affine=False, quant_config=None, prefix="norm_out" ) self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, prior_token_id: torch.Tensor, prior_token_drop: torch.Tensor, timestep: torch.LongTensor, target_size: torch.Tensor, crop_coords: torch.Tensor, attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, attention_mask: torch.Tensor | None = None, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, kv_cache: GlmImageKVCache | None = None, ) -> torch.Tensor | Transformer2DModelOutput: """ Forward pass of the GLM-Image Transformer. Args: hidden_states: Input latent tensor of shape [B, C, H, W]. encoder_hidden_states: Text embeddings of shape [B, S, D]. prior_token_id: Prior VQ token IDs. prior_token_drop: Mask for dropping prior tokens (CFG). timestep: Diffusion timestep. target_size: Target image size for conditioning. crop_coords: Crop coordinates for conditioning. attention_kwargs: Additional attention arguments. return_dict: Whether to return a dataclass. attention_mask: Optional attention mask for text tokens. image_rotary_emb: Pre-computed rotary embeddings. kv_cache: Optional KV cache for image editing. When provided, the cache's mode determines behavior: - WRITE: Store KV from condition images - READ: Use cached KV during generation - SKIP: No caching (same as None) Returns: Output tensor or Transformer2DModelOutput. """ batch_size, num_channels, height, width = hidden_states.shape # Get KV cache mode kv_cache_mode = kv_cache.mode if kv_cache is not None else None # Set SP context if enabled sp_size = self.parallel_config.sequence_parallel_size if sp_size is not None and sp_size > 1: get_forward_context().split_text_embed_in_sp = False # Text embedding projection encoder_hidden_states = self.glyph_projector(encoder_hidden_states) # Prior embedding with dropout prior_embedding = self.prior_token_embedding(prior_token_id) prior_embedding[prior_token_drop] *= 0.0 prior_hidden_states = self.prior_projector(prior_embedding) # 1. Prepare hidden_states and RoPE via GlmImagePrepare module # _sp_plan will shard hidden_states and RoPE together via split_output=True hidden_states, rope_cos, rope_sin, post_patch_height_t, post_patch_width_t = self.prepare( hidden_states, prior_hidden_states ) image_rotary_emb = (rope_cos, rope_sin) post_patch_height = int(post_patch_height_t.item()) post_patch_width = int(post_patch_width_t.item()) # Timestep conditioning temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) # Create padding mask for SP if needed (after _sp_plan hooks have run) hidden_states_mask = None if sp_size is not None and sp_size > 1: from vllm_omni.diffusion.forward_context import is_forward_context_available if is_forward_context_available(): ctx = get_forward_context() if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size hidden_states_mask = torch.ones( batch_size, img_padded_seq_len, dtype=torch.bool, device=hidden_states.device, ) hidden_states_mask[:, ctx.sp_original_seq_len :] = False if hidden_states_mask.all(): hidden_states_mask = None # 2. Transformer blocks for layer_idx, block in enumerate(self.transformer_blocks): # Get layer-specific KV cache if available layer_kv_cache = kv_cache[layer_idx] if kv_cache is not None else None hidden_states, encoder_hidden_states = block( hidden_states, encoder_hidden_states, temb, image_rotary_emb, attention_mask, attention_kwargs, kv_cache=layer_kv_cache, kv_cache_mode=kv_cache_mode, hidden_states_mask=hidden_states_mask, ) # 3. Output norm & projection # _sp_plan will gather hidden_states via proj_out hook hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.proj_out(hidden_states) # 4. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W] p = self.patch_size hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """ Load weights from pretrained checkpoint. This method handles the mapping from diffusers weight names to vllm-omni weight names, especially for fused QKV projections. """ stacked_params_mapping = [ # Fused QKV projection: to_q, to_k, to_v -> to_qkv (".to_qkv", ".to_q", "q"), (".to_qkv", ".to_k", "k"), (".to_qkv", ".to_v", "v"), ] # Expose packed shard mappings for LoRA handling of fused projections. self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) # Also include buffers (for any beta/eps parameters) for name, buffer in self.named_buffers(): params_dict[name] = buffer loaded_params: set[str] = set() for name, loaded_weight in weights: # Handle fused QKV projections for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue # Map diffusers name to vllm-omni name name = name.replace(weight_name, param_name) if name not in params_dict: logger.warning(f"Skipping weight {name} - not found in model") break param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight, shard_id) break else: # Standard weight loading (not fused) if name not in params_dict: logger.warning(f"Skipping weight {name} - not found in model") continue 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 def create_kv_cache(self) -> GlmImageKVCache: """ Create a KV cache for image editing. Returns a new GlmImageKVCache instance sized for this model's number of transformer layers. Use this for image editing workflows. Example: kv_cache = transformer.create_kv_cache() kv_cache.set_mode("write") transformer(condition_image, kv_cache=kv_cache) kv_cache.set_mode("read") for t in timesteps: transformer(noisy_target, kv_cache=kv_cache) kv_cache.clear() Returns: GlmImageKVCache instance with correct number of layers. """ return GlmImageKVCache(num_layers=len(self.transformer_blocks)) @property def num_layers(self) -> int: """Return number of transformer layers.""" return len(self.transformer_blocks) @property def dtype(self) -> torch.dtype: """Return dtype of model parameters.""" return next(self.parameters()).dtype