# Copyright 2025 The ZhipuAI Inc. team and HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections.abc import Callable from itertools import accumulate import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from huggingface_hub.dataclasses import strict from torch.nn import LayerNorm from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig, SubConfigSpec from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack from ...utils import ( TransformersKwargs, auto_docstring, logging, torch_compilable_check, ) from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ...vision_utils import get_vision_attention_seqlens, get_vision_position_ids from ..glm4.modeling_glm4 import Glm4MLP, Glm4RMSNorm, eager_attention_forward from ..qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionPatchEmbed, Qwen2_5_VLCausalLMOutputWithPast, Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLMLP, Qwen2_5_VLModelOutputWithPast, Qwen2_5_VLPreTrainedModel, Qwen2_5_VLRotaryEmbedding, Qwen2_5_VLTextModel, Qwen2_5_VLVisionAttention, Qwen2_5_VLVisionBlock, Qwen2_5_VLVisionRotaryEmbedding, ) from ..qwen2_vl.modeling_qwen2_vl import Qwen2VLModel from ..qwen2_vl.processing_qwen2_vl import ( Qwen2VLProcessor, Qwen2VLProcessorKwargs, ) logger = logging.get_logger(__name__) @auto_docstring(checkpoint="zai-org/GLM-4.1V-9B-Thinking") @strict class Glm4vVisionConfig(PreTrainedConfig): r""" out_hidden_size (`int`, *optional*, defaults to 4096): The output hidden size of the vision model. Example: ```python >>> from transformers import Glm4vVisionConfig, Glm4vVisionModel >>> # Initializing a Glm4vVisionConfig GLM-4.1V-9B style configuration >>> configuration = Glm4vVisionConfig() >>> # Initializing a model (with random weights) from the GLM-4.1V-9B configuration >>> model = Glm4vVisionModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "glm4v_vision" base_config_key = "vision_config" default_rope_type = "axial" attribute_map = {"num_attention_heads": "num_heads"} depth: int = 24 hidden_size: int = 1536 hidden_act: str = "silu" attention_bias: bool = False attention_dropout: float | int = 0.0 num_heads: int = 12 in_channels: int = 3 image_size: int | list[int] | tuple[int, int] = 336 patch_size: int | list[int] | tuple[int, int] = 14 rms_norm_eps: float = 1e-05 spatial_merge_size: int = 2 temporal_patch_size: int | list[int] | tuple[int, int] = 2 out_hidden_size: int = 4096 intermediate_size: int = 13696 initializer_range: float = 0.02 rope_parameters: dict | None = None @auto_docstring(checkpoint="zai-org/GLM-4.1V-9B-Thinking") @strict class Glm4vTextConfig(PreTrainedConfig): r""" Example: ```python >>> from transformers import Glm4vTextModel, Glm4vConfig >>> # Initializing a GLM-4.1V style configuration >>> configuration = Glm4vConfig() >>> # Initializing a model from the GLM-4.1V style configuration >>> model = Glm4vTextModel(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "glm4v_text" base_config_key = "text_config" keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Glm4v` base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } ignore_keys_at_rope_validation = {"mrope_section"} vocab_size: int = 151552 hidden_size: int = 4096 intermediate_size: int = 13696 num_hidden_layers: int = 40 num_attention_heads: int = 32 num_key_value_heads: int | None = 2 hidden_act: str = "silu" max_position_embeddings: int = 32768 initializer_range: float = 0.02 rms_norm_eps: float = 1e-05 use_cache: bool = True attention_dropout: float | int = 0.0 rope_parameters: RopeParameters | dict | None = None pad_token_id: int | None = None def __post_init__(self, **kwargs): if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads super().__post_init__(**kwargs) @auto_docstring(checkpoint="zai-org/GLM-4.1V-9B-Thinking") @strict class Glm4vConfig(PreTrainedConfig): r""" image_start_token_id (`int`, *optional*, defaults to 151339): The image start token index to encode the start of image. image_end_token_id (`int`, *optional*, defaults to 151340): The image end token index to encode the end of image. video_start_token_id (`int`, *optional*, defaults to 151341): The video start token index to encode the start of video. video_end_token_id (`int`, *optional*, defaults to 151342): The video end token index to encode the end of video. ```python >>> from transformers import Glm4vForConditionalGeneration, Glm4vConfig >>> # Initializing a GLM-4.1V style configuration >>> configuration = Glm4vConfig() >>> # Initializing a model from the GLM-4.1V style configuration >>> model = Glm4vForConditionalGeneration(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "glm4v" sub_configs_defaults = { "vision_config": SubConfigSpec(config_class=Glm4vVisionConfig), "text_config": SubConfigSpec(config_class=Glm4vTextConfig), } keys_to_ignore_at_inference = ["past_key_values"] text_config: dict | PreTrainedConfig | None = None vision_config: dict | PreTrainedConfig | None = None image_token_id: int = 151343 video_token_id: int = 151344 image_start_token_id: int = 151339 image_end_token_id: int = 151340 video_start_token_id: int = 151341 video_end_token_id: int = 151342 tie_word_embeddings: bool = False def __post_init__(self, **kwargs): super().__post_init__(**kwargs) # BC: pre-v5 saves placed `tie_word_embeddings` inside text_config. Forward it to the outer # config (where v5's tying logic looks) when the root value is the default. Checked after # text_config init so it also covers a text config passed as an already-initialized instance. if not self.tie_word_embeddings and getattr(self.text_config, "tie_word_embeddings", False): self.tie_word_embeddings = True # Will be used for both Text and Vision modalities class Glm4vRMSNorm(Glm4RMSNorm): pass class Glm4VisionMlp(Qwen2_5_VLMLP): def __init__(self, config, bias: bool = False): super().__init__(config, bias) self.intermediate_size = config.out_hidden_size class Glm4vVisionPatchEmbed(Qwen2_5_VisionPatchEmbed): def __init__(self, config: Glm4vVisionConfig) -> None: nn.Module.__init__(self) self.patch_size = config.patch_size self.temporal_patch_size = config.temporal_patch_size self.in_channels = config.in_channels self.embed_dim = config.hidden_size kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size) class Glm4vVisionRotaryEmbedding(Qwen2_5_VLVisionRotaryEmbedding): pass class Glm4vVisionPatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, hidden_act: str, bias: bool = False) -> None: super().__init__() self.proj = nn.Linear(dim, dim, bias=bias) self.post_projection_norm = LayerNorm(dim) self.gate_proj = nn.Linear(dim, context_dim, bias=bias) self.up_proj = nn.Linear(dim, context_dim, bias=bias) self.down_proj = nn.Linear(context_dim, dim, bias=bias) self.act1 = nn.GELU() self.act_fn = ACT2FN[hidden_act] def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: hidden_state = self.proj(hidden_state) hidden_state = self.act1(self.post_projection_norm(hidden_state)) return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) class Glm4vVisionEmbeddings(nn.Module): def __init__(self, config: Glm4vVisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.image_size = config.image_size self.patch_size = config.patch_size self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) self.interpolated_method = "bicubic" def forward(self, embeddings, lengths, image_shapes, h_coords, w_coords) -> torch.Tensor: """ Forward pass with integrated position encoding adaptation using 2D interpolation. Args: embeddings: Input embeddings tensor lengths (torch.Tensor): Sequence lengths for each image in the batch. image_shapes (torch.Tensor): Tensor of shape [batch_size, 3] representing the image shapes (t, h, w). h_coords (torch.Tensor): Tensor of shape [total_seq] representing the h coordinate for each patch. w_coords (torch.Tensor): Tensor of shape [total_seq] representing the w coordinate for each patch. Returns: torch.Tensor: Embeddings with adapted position encoding added. """ # Get position embedding parameters pos_embed_weight = self.position_embedding.weight hidden_size = pos_embed_weight.shape[1] device = pos_embed_weight.device # Convert inputs to tensors if needed if isinstance(lengths, list): lengths = torch.tensor(lengths, device=device, dtype=torch.long) # Prepare 2D position embedding orig_size_sq = pos_embed_weight.shape[0] orig_size = int(orig_size_sq**0.5) pos_embed_2d = ( pos_embed_weight.view(orig_size, orig_size, hidden_size) .permute(2, 0, 1) .unsqueeze(0) .to(device=device, dtype=torch.float32) ) # Calculate target dimensions for each patch num_tokens = embeddings.shape[0] token_positions = torch.arange(num_tokens, device=embeddings.device) seq_ids = (token_positions.unsqueeze(0) >= lengths.cumsum(0).unsqueeze(1)).sum(0) target_h = image_shapes[seq_ids, 1].to(dtype=torch.float32) target_w = image_shapes[seq_ids, 2].to(dtype=torch.float32) # Normalize coordinates to [-1, 1] range for grid_sample norm_w = ((w_coords + 0.5) / target_w) * 2 - 1 norm_h = ((h_coords + 0.5) / target_h) * 2 - 1 # Create sampling grid grid = torch.stack((norm_w, norm_h), dim=-1).unsqueeze(0).unsqueeze(2) # Perform bicubic interpolation interpolated_embed_fp32 = F.grid_sample( pos_embed_2d, grid, mode=self.interpolated_method, align_corners=False, padding_mode="border" ) # Reshape and convert back to original dtype adapted_pos_embed_fp32 = interpolated_embed_fp32.squeeze(0).squeeze(-1).permute(1, 0) adapted_pos_embed = adapted_pos_embed_fp32.to(pos_embed_weight.dtype).to(embeddings.device) # Add adapted position encoding to embeddings embeddings = embeddings + adapted_pos_embed return embeddings class Glm4vVisionAttention(Qwen2_5_VLVisionAttention): def __init__(self, config: Glm4vVisionConfig) -> None: super().__init__(config) self.attention_dropout = config.attention_dropout self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias) self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) class Glm4vVisionBlock(Qwen2_5_VLVisionBlock): def __init__(self, config) -> None: super().__init__(config) self.norm1 = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.norm2 = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attn = Glm4vVisionAttention(config) self.mlp = Glm4VisionMlp(config, bias=False) class Glm4vTextRotaryEmbedding(Qwen2_5_VLRotaryEmbedding): def __init__(self, config: Glm4vTextConfig): super().__init__() self.mrope_section = config.rope_parameters.get("mrope_section", [8, 12, 12]) def compute_default_rope_parameters(config: Glm4vConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads dim = int(head_dim * partial_rotary_factor) attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) return inv_freq.to(device), attention_factor def recomposition_frequencies(self, freq): """ Recompose the frequencies into the final spatial layout used per each grid. """ freq = torch.cat([m[i % 3] for i, m in enumerate(freq.split(self.mrope_section, dim=-1))], dim=-1) return freq.repeat_interleave(2, dim=-1) def rotate_half_llm(x): """Rotates half the hidden dims of the input.""" x1 = x[..., 0::2] x2 = x[..., 1::2] return torch.stack((-x2, x1), dim=-1).flatten(-2) def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) # Keep half or full tensor for later concatenation rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] # Apply rotary embeddings on the first half or full tensor q_embed = (q_rot * cos) + (rotate_half_llm(q_rot) * sin) k_embed = (k_rot * cos) + (rotate_half_llm(k_rot) * sin) # Concatenate back to full shape q_embed = torch.cat([q_embed, q_pass], dim=-1) k_embed = torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed class Glm4vTextAttention(nn.Module): """ Multi-headed attention from 'Attention Is All You Need' paper. and "Generating Long Sequences with Sparse Transformers". """ def __init__(self, config: Glm4vTextConfig, layer_idx: int | None = None): super().__init__() self.config = config self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.is_causal = True self.attention_dropout = config.attention_dropout self.rope_parameters = config.rope_parameters self.scaling = self.head_dim**-0.5 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class Glm4vTextMLP(Glm4MLP): pass class Glm4vTextDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: Glm4vTextConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = Glm4vTextAttention(config, layer_idx) self.mlp = Glm4vTextMLP(config) self.input_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_self_attn_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_mlp_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @auto_docstring def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, **kwargs, ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, **kwargs, ) hidden_states = self.post_self_attn_layernorm(hidden_states) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.post_mlp_layernorm(hidden_states) hidden_states = residual + hidden_states return hidden_states class Glm4vModelOutputWithPast(Qwen2_5_VLModelOutputWithPast): pass class Glm4vPreTrainedModel(Qwen2_5_VLPreTrainedModel): _no_split_modules = ["Glm4vTextDecoderLayer", "Glm4vVisionBlock"] class Glm4vVisionModel(Glm4vPreTrainedModel): config: Glm4vVisionConfig input_modalities = ("image", "video") _no_split_modules = ["Glm4vVisionBlock"] _can_record_outputs = { "hidden_states": Glm4vVisionBlock, "attentions": Glm4vVisionAttention, } def __init__(self, config) -> None: super().__init__(config) self.spatial_merge_size = config.spatial_merge_size self.patch_size = config.patch_size self.embeddings = Glm4vVisionEmbeddings(config) self.patch_embed = Glm4vVisionPatchEmbed(config) self.rotary_pos_emb = Glm4vVisionRotaryEmbedding(config) self.blocks = nn.ModuleList([Glm4vVisionBlock(config) for _ in range(config.depth)]) self.merger = Glm4vVisionPatchMerger( dim=config.out_hidden_size, context_dim=config.intermediate_size, hidden_act=config.hidden_act ) self.post_conv_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.downsample = nn.Conv2d( in_channels=config.hidden_size, out_channels=config.out_hidden_size, kernel_size=config.spatial_merge_size, stride=config.spatial_merge_size, ) self.post_layernorm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] ) -> tuple | BaseModelOutputWithPooling: r""" hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): The final hidden states of the model. grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): The temporal, height and width of feature shape of each image in LLM. Returns: `torch.Tensor`: hidden_states. """ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs) cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs) hidden_states = self.patch_embed(hidden_states) hidden_states = self.post_conv_layernorm(hidden_states) position_embeddings = self.rotary_pos_emb(hidden_states, position_ids) seqlens = cu_seqlens[1:] - cu_seqlens[:-1] hidden_states = self.embeddings( hidden_states, seqlens, grid_thw, position_ids[:, 0].to(hidden_states.device), position_ids[:, 1].to(hidden_states.device), ) for blk in self.blocks: hidden_states = blk( hidden_states, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.post_layernorm(hidden_states) hidden_states = hidden_states.view( -1, self.spatial_merge_size, self.spatial_merge_size, hidden_states.shape[-1] ) hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size) merged_hidden_states = self.merger(hidden_states) return BaseModelOutputWithPooling( last_hidden_state=hidden_states, pooler_output=merged_hidden_states, ) class Glm4vTextModel(Qwen2_5_VLTextModel): _can_record_outputs = { "hidden_states": Glm4vTextDecoderLayer, "attentions": Glm4vTextAttention, } def __init__(self, config: Glm4vTextConfig): super().__init__(config) self.layers = nn.ModuleList( [Glm4vTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Glm4vRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Glm4vTextRotaryEmbedding(config=config) del self._attn_implementation del self.has_sliding_layers @auto_docstring @merge_with_config_defaults @capture_outputs def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") # torch.jit.trace() doesn't support cache objects in the output if use_cache and past_key_values is None and not torch.jit.is_tracing(): past_key_values = DynamicCache(config=self.config) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) # the hard coded `3` is for temporal, height and width. if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) elif position_ids.ndim == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions # where each dim indicates visual spatial positions for temporal/height/width grids. # There are two scenarios when FA2-like packed masking might be activated. # 1. User specifically passed packed `position_ids` and no attention mask. # In this case we expect the user to create correct position ids for all 3 grids # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len] # 2. User runs forward with no attention mask and no position ids. In this case, position ids # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation` if position_ids.ndim == 3 and position_ids.shape[0] == 4: text_position_ids = position_ids[0] position_ids = position_ids[1:] else: # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids text_position_ids = None mask_kwargs = { "config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "past_key_values": past_key_values, "position_ids": text_position_ids, } # Create the masks causal_mask = create_causal_mask(**mask_kwargs) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) for decoder_layer in self.layers: layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=text_position_ids, past_key_values=past_key_values, position_embeddings=position_embeddings, **kwargs, ) hidden_states = layer_outputs hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) class Glm4vModel(Qwen2VLModel): _no_split_modules = ["Glm4vTextDecoderLayer", "Glm4vVisionBlock"] def __init__(self, config): super().__init__(config) self.visual = Glm4vVisionModel._from_config(config.vision_config) def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: pixel_values_videos = pixel_values_videos.type(self.visual.dtype) # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames t = video_grid_thw[:, 0] hw = video_grid_thw[:, 1:] # repeat each (h,w) row `t` times flattened_hw = torch.repeat_interleave(hw, t, dim=0) prefix_ones = video_grid_thw.new_ones(flattened_hw.shape[0], 1) flattened_video_grid_thw = torch.cat([prefix_ones, flattened_hw], dim=1) vision_outputs = self.visual( pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs ) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() vision_outputs.pooler_output = torch.split(vision_outputs.pooler_output, split_sizes) return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor | None = None, video_features: torch.FloatTensor | None = None, ): """ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is equal to the length of multimodal features. If the lengths are different, an error is raised. """ if input_ids is None: special_image_mask = inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_image_mask = special_image_mask.all(-1) special_video_mask = inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_video_mask = special_video_mask.all(-1) else: # GLM-4.1V and GLM-4.5V special_video_mask is special_image_mask special_image_mask = input_ids == self.config.image_token_id special_video_mask = input_ids == self.config.image_token_id n_image_tokens = special_image_mask.sum() special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device) if image_features is not None: torch_compilable_check( n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(), f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}", ) n_video_tokens = special_video_mask.sum() special_video_mask = special_video_mask.unsqueeze(-1).to(inputs_embeds.device) if video_features is not None: torch_compilable_check( n_video_tokens * inputs_embeds.shape[-1] == video_features.numel(), f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}", ) return special_image_mask, special_video_mask def get_rope_index( self, video_grid_thw: torch.LongTensor | None = None, **super_kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ Difference from Qwen2VL/Qwen2.5VL's get_rope_index: - GLM4V uses timestamps to separate each video frame, so the video_grid_thw should also be split too. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. Returns: position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ # Separate video grid thw into multiple grids because timestamps are used to separate videos. if video_grid_thw is not None: video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) video_grid_thw[:, 0] = 1 return super().get_rope_index(video_grid_thw=video_grid_thw, **super_kwargs) def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, mm_token_type_ids: torch.IntTensor | None = None, mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm4vModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None: raise ValueError( "You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time" ) if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {} if mm_encoder_outputs.get("image") is None and pixel_values is not None: mm_encoder_outputs["image"] = self.get_image_features( pixel_values, image_grid_thw, return_dict=True, **kwargs ) if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None: mm_encoder_outputs["video"] = self.get_video_features( pixel_values_videos, video_grid_thw, return_dict=True, **kwargs ) if mm_encoder_outputs.get("image") is not None: image_embeds = mm_encoder_outputs["image"].pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if mm_encoder_outputs.get("video") is not None: video_embeds = mm_encoder_outputs["video"].pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) if position_ids is None: position_ids = self.compute_3d_position_ids( input_ids=input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( input_ids=None, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs, ) return Glm4vModelOutputWithPast( **outputs, rope_deltas=self.rope_deltas, ) class Glm4vCausalLMOutputWithPast(Qwen2_5_VLCausalLMOutputWithPast): pass class Glm4vForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, mm_token_type_ids: torch.IntTensor | None = None, logits_to_keep: int | torch.Tensor = 0, mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm4vCausalLMOutputWithPast: r""" Example: ```python >>> from PIL import Image >>> from huggingface_hub.utils import httpx >>> from io import BytesIO >>> from transformers import AutoProcessor, Glm4vForConditionalGeneration >>> model = Glm4vForConditionalGeneration.from_pretrained("zai-org/GLM-4.1V-9B-Thinking") >>> processor = AutoProcessor.from_pretrained("zai-org/GLM-4.1V-9B-Thinking") >>> messages = [ { "role": "user", "content": [ {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"}, {"type": "text", "text": "What is shown in this image?"}, ], }, ] >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" >>> with httpx.stream("GET", url) as response: ... image = Image.open(BytesIO(response.read())) >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos]) >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..." ```""" outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, mm_encoder_outputs=mm_encoder_outputs, **kwargs, ) hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs ) return Glm4vCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, ) def _expand_multimodal_outputs( self, input_ids: torch.LongTensor, mm_encoder_output: dict, expand_size: int = 1, inputs_embeds: torch.LongTensor | None = None, ) -> dict[str, dict]: # override -> model uses the same config key for image and video placeholders def repeat_tensor_or_list(inputs: list | torch.Tensor, repeat_times: int): # Tensor of size [bs, seqlen, dim] where `bs` is number of images in this text sample # Each text can have 1+ images associated with it # Inteleaving on first dim does the same thing as `input_ids.repeat_interlave` in leading batch dim! if isinstance(inputs, torch.Tensor): return inputs.repeat_interleave(repeat_times, dim=0) else: # List of `bs` length where each entry is a tensor (seqlen, dim) is also repeat interleaved return [beam_entry for entry in inputs for beam_entry in [entry] * repeat_times] for modality in ["image", "video"]: modalily_outputs = mm_encoder_output.get(modality) if modalily_outputs is None or getattr(modalily_outputs, "pooler_output", None) is None: continue # 1. compute cumulative number of placehlder tokens per sample and per each encoded mm-data token_id_key = f"{modality}_token_id" if (input_ids is None or input_ids.numel() == 0) and inputs_embeds is not None: special_image_mask = inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) num_image_tokens_in_text = special_image_mask.all(-1).sum(-1) else: num_image_tokens_in_text = (input_ids == getattr(self.config, token_id_key)).sum(-1) num_image_tokens_in_vision = [len(out) for out in modalily_outputs.pooler_output] num_image_tokens_in_text = list(accumulate(num_image_tokens_in_text)) num_image_tokens_in_vision = list(accumulate(num_image_tokens_in_vision)) # 2. Find offsets to split encoder output into separate groups per text. In a single batch # we might get a text with single image and another with two images, so the most reliable # way to split dynamic-sized images is by checking number of placeholders and encoder output lengths! offsets = [0] + [ i + 1 for i, num in enumerate(num_image_tokens_in_vision) if num in num_image_tokens_in_text ] is_split_images = not isinstance(modalily_outputs.pooler_output, torch.Tensor) modalily_outputs.pooler_output = [ out for start, end in zip(offsets[:-1], offsets[1:]) for out in repeat_tensor_or_list(modalily_outputs.pooler_output[start:end], expand_size) ] # 3. if `pooler_output` was a tensor, cat it back to follow model expectations if not is_split_images: modalily_outputs.pooler_output = torch.stack(modalily_outputs.pooler_output, dim=0) return mm_encoder_output class Glm4vProcessorKwargs(Qwen2VLProcessorKwargs): _defaults = { "text_kwargs": { "padding": False, "return_token_type_ids": False, "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } class Glm4vProcessor(Qwen2VLProcessor): def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs): super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) self.image_token = "<|image|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token self.video_token = "<|video|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token self.video_start_id = tokenizer.convert_tokens_to_ids("<|begin_of_video|>") self.video_end_id = tokenizer.convert_tokens_to_ids("<|end_of_video|>") def replace_video_token(self, video_inputs: dict, video_idx: int, **kwargs) -> str: merge_length = self.video_processor.merge_size**2 num_frames = video_inputs["video_grid_thw"][video_idx][0] num_image_tokens = video_inputs["video_grid_thw"][video_idx].prod() // merge_length // num_frames metadata = video_inputs["video_metadata"][video_idx] video_structure = "" if metadata.fps is None: logger.warning_once( "GLM4V requires frame timestamps to construct prompts, but the `fps` of the input video could not be inferred. " "Probably `video_metadata` was missing from inputs and you passed pre-sampled frames. " "Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results." ) metadata.fps = 24 if metadata.fps is None else metadata.fps timestamps = metadata.timestamps[::2] # mrope unique_timestamps = [] for idx in range(0, len(timestamps)): unique_timestamps.append(timestamps[idx]) selected_timestamps = unique_timestamps[:num_frames] while len(selected_timestamps) < num_frames: selected_timestamps.append(selected_timestamps[-1] if selected_timestamps else 0) for frame_idx in range(num_frames): timestamp_sec = selected_timestamps[frame_idx] frame_structure = self.replace_frame_token_id(timestamp_sec, num_image_tokens=num_image_tokens) video_structure += frame_structure return video_structure def create_mm_token_type_ids(self, input_ids: list) -> list[list[int]]: # We have to iterate for each list separately because inputs # might be non-padded lists and we can't cast numpy on that! # Then cast numpy as each input for faster indexing mm_token_type_ids = [] for input in input_ids: array_ids = np.array(input) mm_token_types = np.zeros_like(input) # Replace 0 -> 2 only inside video segments because GLM4v # uses the same special token to denote images and video # Otherwise replace 0 -> 1 for image modality starts = np.cumsum(array_ids == self.video_start_id, axis=0) ends = np.cumsum(array_ids == self.video_end_id, axis=0) is_video_modality = starts > ends mm_token_types[(array_ids == self.image_token_id) & is_video_modality] = 2 mm_token_types[(array_ids == self.image_token_id) & (~is_video_modality)] = 1 mm_token_type_ids.append(mm_token_types.tolist()) return mm_token_type_ids def replace_frame_token_id(self, timestamp_sec, num_image_tokens: int = 1): return f"<|begin_of_image|>{self.image_token * num_image_tokens}<|end_of_image|>{int(timestamp_sec)}" __all__ = [ "Glm4vConfig", "Glm4vTextConfig", "Glm4vVisionConfig", "Glm4vForConditionalGeneration", "Glm4vModel", "Glm4vPreTrainedModel", "Glm4vProcessor", "Glm4vTextModel", "Glm4vVisionModel", ]