From 393cb1aadf5c994f4c48aa521655596cf24dbace Mon Sep 17 00:00:00 2001 From: trueoneplusone <78432179+trueoneplusone@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:52:54 +0800 Subject: [PATCH] Add native CogVLM2-Video support --- docs/source/en/_toctree.yml | 4 +- docs/source/en/model_doc/cogvlm2.md | 95 +++ src/transformers/models/auto/auto_mappings.py | 9 +- src/transformers/models/auto/modeling_auto.py | 2 + src/transformers/models/cogvlm2/__init__.py | 29 + .../models/cogvlm2/configuration_cogvlm2.py | 116 ++++ .../models/cogvlm2/modeling_cogvlm2.py | 631 +++++++++++++++++ .../models/cogvlm2/modular_cogvlm2.py | 655 ++++++++++++++++++ .../models/cogvlm2/processing_cogvlm2.py | 236 +++++++ .../cogvlm2/video_processing_cogvlm2.py | 95 +++ tests/models/cogvlm2/__init__.py | 0 tests/models/cogvlm2/test_modeling_cogvlm2.py | 128 ++++ .../models/cogvlm2/test_processing_cogvlm2.py | 110 +++ .../cogvlm2/test_video_processing_cogvlm2.py | 38 + 14 files changed, 2145 insertions(+), 3 deletions(-) create mode 100644 docs/source/en/model_doc/cogvlm2.md create mode 100644 src/transformers/models/cogvlm2/__init__.py create mode 100644 src/transformers/models/cogvlm2/configuration_cogvlm2.py create mode 100644 src/transformers/models/cogvlm2/modeling_cogvlm2.py create mode 100644 src/transformers/models/cogvlm2/modular_cogvlm2.py create mode 100644 src/transformers/models/cogvlm2/processing_cogvlm2.py create mode 100644 src/transformers/models/cogvlm2/video_processing_cogvlm2.py create mode 100644 tests/models/cogvlm2/__init__.py create mode 100644 tests/models/cogvlm2/test_modeling_cogvlm2.py create mode 100644 tests/models/cogvlm2/test_processing_cogvlm2.py create mode 100644 tests/models/cogvlm2/test_video_processing_cogvlm2.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index e069c7966028..8d9491efff86 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -832,7 +832,7 @@ - local: model_doc/nomic_bert title: NomicBERT - local: model_doc/nystromformer - title: Nystrâmformer + title: Nystr?mformer - local: model_doc/olmo title: OLMo - local: model_doc/olmo2 @@ -1281,6 +1281,8 @@ title: CLIPSeg - local: model_doc/clvp title: CLVP + - local: model_doc/cogvlm2 + title: CogVLM2 - local: model_doc/cwm title: Code World Model (CWM) - local: model_doc/cohere2_vision diff --git a/docs/source/en/model_doc/cogvlm2.md b/docs/source/en/model_doc/cogvlm2.md new file mode 100644 index 000000000000..5964396d4c9a --- /dev/null +++ b/docs/source/en/model_doc/cogvlm2.md @@ -0,0 +1,95 @@ + + +# CogVLM2 + +## Overview + +CogVLM2-Video is a video-language model from Zhipu AI / THUDM. The released +CogVLM2-Video Llama 3 checkpoint combines an EVA2-CLIP visual encoder with a +Llama 3-style causal language model. + +Each video frame is encoded into 64 spatial visual tokens. A beginning-of-image +and end-of-image embedding are added around those tokens, producing 66 visual +tokens per frame. CogVLM2 uses a compressed position-id scheme for visual spans +so that long visual token runs do not consume one language position per token. + +The native Transformers implementation supports the modern cache and generation +APIs and does not depend on the original Triton rotary kernel. + +## Inference + +Once the checkpoint metadata has been updated for native Transformers support, +CogVLM2-Video can be loaded without remote code: + +```python +from transformers import AutoModelForCausalLM, AutoProcessor + + +model_id = "zai-org/cogvlm2-video-llama3-chat" + +model = AutoModelForCausalLM.from_pretrained( + model_id, + dtype="auto", + device_map="auto", +) +processor = AutoProcessor.from_pretrained(model_id) + +inputs = processor( + text="What happens in this video?", + videos="path/to/video.mp4", + return_tensors="pt", +).to(model.device) + +output_ids = model.generate(**inputs, max_new_tokens=128) +print(processor.tokenizer.decode(output_ids[0], skip_special_tokens=True)) +``` + +The original checkpoint was published with remote-code metadata and without a +`model_type` field. Existing Hub repositories need their configuration and +processor metadata updated to use native loading after this model is released. + +## CogVLM2Config + +[[autodoc]] CogVLM2Config + +## CogVLM2VisionConfig + +[[autodoc]] CogVLM2VisionConfig + +## CogVLM2VideoProcessor + +[[autodoc]] CogVLM2VideoProcessor + - preprocess + +## CogVLM2Processor + +[[autodoc]] CogVLM2Processor + - __call__ + +## CogVLM2VisionModel + +[[autodoc]] CogVLM2VisionModel + - forward + +## CogVLM2Model + +[[autodoc]] CogVLM2Model + - forward + - get_video_features + +## CogVLM2ForConditionalGeneration + +[[autodoc]] CogVLM2ForConditionalGeneration + - forward + - prepare_inputs_for_generation diff --git a/src/transformers/models/auto/auto_mappings.py b/src/transformers/models/auto/auto_mappings.py index 5778b7aa3ba2..223e5040b391 100644 --- a/src/transformers/models/auto/auto_mappings.py +++ b/src/transformers/models/auto/auto_mappings.py @@ -1,8 +1,8 @@ -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# ???????????????????????????????????????????????????????????????????????????????????????????????? # This file was automatically generated from existing config files and their `model_type`s. Do NOT edit this file # manually as any edits will be overwritten by auto-generation of the file. If any change should be done, # please add the correct `cls.model_type` in your config class and run `python utils/check_auto.py --fix_and_overwrite`. -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# ???????????????????????????????????????????????????????????????????????????????????????????????? # Copyright 2026 The HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -97,6 +97,8 @@ ("clvp_decoder", "ClvpDecoderConfig"), ("clvp_encoder", "ClvpEncoderConfig"), ("codegen", "CodeGenConfig"), + ("cogvlm2", "CogVLM2Config"), + ("cogvlm2_vision", "CogVLM2VisionConfig"), ("cohere", "CohereConfig"), ("cohere2", "Cohere2Config"), ("cohere2_moe", "Cohere2MoeConfig"), @@ -976,6 +978,7 @@ ("tipsv2_vision_model", "tipsv2"), ("unispeech-sat", "unispeech_sat"), ("uvdoc_backbone", "uvdoc"), + ("cogvlm2_vision", "cogvlm2"), ("video_llama_3_vision", "video_llama_3"), ("videoprism_text_model", "videoprism"), ("videoprism_vision_model", "videoprism"), @@ -996,6 +999,7 @@ VIDEO_PROCESSOR_MAPPING_NAMES = OrderedDict( [ + ("cogvlm2", "CogVLM2VideoProcessor"), ("cohere_compass", "CohereCompassVideoProcessor"), ("cosmos3_edge", "Cosmos3EdgeVideoProcessor"), ("ernie4_5_vl_moe", "Ernie4_5_VLMoeVideoProcessor"), @@ -1082,6 +1086,7 @@ ("clip", "CLIPProcessor"), ("clipseg", "CLIPSegProcessor"), ("clvp", "ClvpProcessor"), + ("cogvlm2", "CogVLM2Processor"), ("cohere2_vision", "Cohere2VisionProcessor"), ("cohere_asr", "CohereAsrProcessor"), ("cohere_compass", "CohereCompassProcessor"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 29924dae5814..741eccf48825 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -91,6 +91,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("clipseg", "CLIPSegModel"), ("clvp", "ClvpModelForConditionalGeneration"), ("codegen", "CodeGenModel"), + ("cogvlm2", "CogVLM2Model"), ("cohere", "CohereModel"), ("cohere2", "Cohere2Model"), ("cohere2_moe", "Cohere2MoeModel"), @@ -723,6 +724,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("blt", "BltForCausalLM"), ("camembert", "CamembertForCausalLM"), ("codegen", "CodeGenForCausalLM"), + ("cogvlm2", "CogVLM2ForConditionalGeneration"), ("cohere", "CohereForCausalLM"), ("cohere2", "Cohere2ForCausalLM"), ("cohere2_moe", "Cohere2MoeForCausalLM"), diff --git a/src/transformers/models/cogvlm2/__init__.py b/src/transformers/models/cogvlm2/__init__.py new file mode 100644 index 000000000000..f1bf36f8c697 --- /dev/null +++ b/src/transformers/models/cogvlm2/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2026 The HuggingFace 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_cogvlm2 import * + from .modeling_cogvlm2 import * + from .processing_cogvlm2 import * + from .video_processing_cogvlm2 import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/cogvlm2/configuration_cogvlm2.py b/src/transformers/models/cogvlm2/configuration_cogvlm2.py new file mode 100644 index 000000000000..be695b81729d --- /dev/null +++ b/src/transformers/models/cogvlm2/configuration_cogvlm2.py @@ -0,0 +1,116 @@ +# ???????????????????????????????????????????????????????????????????????????????????????????????? +# This file was automatically generated from src/transformers/models/cogvlm2/modular_cogvlm2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_cogvlm2.py file directly. One of our CI enforces this. +# ???????????????????????????????????????????????????????????????????????????????????????????????? +# Copyright 2026 The HuggingFace 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 ...configuration_utils import PreTrainedConfig + + +class CogVLM2VisionConfig(PreTrainedConfig): + model_type = "cogvlm2_vision" + base_config_key = "vision_config" + + def __init__( + self, + hidden_size=1792, + image_size=224, + in_channels=3, + intermediate_size=15360, + layer_norm_eps=1e-6, + num_heads=16, + num_hidden_layers=63, + num_positions=257, + patch_size=14, + dropout_prob=0.0, + hidden_act="gelu", + **kwargs, + ): + self.hidden_size = hidden_size + self.image_size = image_size + self.in_channels = in_channels + self.intermediate_size = intermediate_size + self.layer_norm_eps = layer_norm_eps + self.num_heads = num_heads + self.num_hidden_layers = num_hidden_layers + self.num_positions = num_positions + self.patch_size = patch_size + self.dropout_prob = dropout_prob + self.hidden_act = hidden_act + super().__init__(**kwargs) + + +class CogVLM2Config(PreTrainedConfig): + model_type = "cogvlm2" + sub_configs = {"vision_config": CogVLM2VisionConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=128256, + hidden_size=4096, + intermediate_size=14336, + num_hidden_layers=32, + num_attention_heads=32, + num_multi_query_heads=8, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-5, + template_version="base", + vision_config=None, + rope_theta=500000.0, + pad_token_id=128002, + bos_token_id=128000, + eos_token_id=128001, + tie_word_embeddings=False, + use_cache=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_multi_query_heads = num_multi_query_heads + self.num_key_value_heads = num_multi_query_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.template_version = template_version + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} + + if vision_config is None: + self.vision_config = CogVLM2VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = CogVLM2VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["CogVLM2VisionConfig", "CogVLM2Config"] diff --git a/src/transformers/models/cogvlm2/modeling_cogvlm2.py b/src/transformers/models/cogvlm2/modeling_cogvlm2.py new file mode 100644 index 000000000000..efb33a52f7be --- /dev/null +++ b/src/transformers/models/cogvlm2/modeling_cogvlm2.py @@ -0,0 +1,631 @@ +# ???????????????????????????????????????????????????????????????????????????????????????????????? +# This file was automatically generated from src/transformers/models/cogvlm2/modular_cogvlm2.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_cogvlm2.py file directly. One of our CI enforces this. +# ???????????????????????????????????????????????????????????????????????????????????????????????? +# Copyright 2026 The HuggingFace 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. + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_kernel_forward_from_hub +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_cogvlm2 import CogVLM2Config + + +class CogVLM2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class CogVLM2PatchEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + self.proj = nn.Conv2d( + config.in_channels, + config.hidden_size, + kernel_size=config.patch_size, + stride=config.patch_size, + ) + self.cls_embedding = nn.Parameter(torch.zeros(1, config.hidden_size)) + self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size) + + def forward(self, pixel_values): + hidden_states = self.proj(pixel_values).flatten(2).transpose(1, 2) + cls_token = self.cls_embedding.expand(hidden_states.shape[0], -1, -1) + hidden_states = torch.cat((cls_token, hidden_states), dim=1) + return hidden_states + self.position_embedding.weight.unsqueeze(0) + + +class CogVLM2VisionAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.hidden_size // config.num_heads + self.query_key_value = nn.Linear(config.hidden_size, config.hidden_size * 3) + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.output_dropout = nn.Dropout(config.dropout_prob) + + def forward(self, hidden_states): + batch_size, seq_len, _ = hidden_states.shape + qkv = self.query_key_value(hidden_states) + qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim) + query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) + attn_output = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1) + return self.output_dropout(self.dense(attn_output)) + + +class CogVLM2VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states): + return self.fc2(self.activation_fn(self.fc1(hidden_states))) + + +class CogVLM2VisionEncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attention = CogVLM2VisionAttention(config) + self.mlp = CogVLM2VisionMLP(config) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states): + residual = hidden_states + hidden_states = residual + self.input_layernorm(self.attention(hidden_states)) + residual = hidden_states + hidden_states = residual + self.post_attention_layernorm(self.mlp(hidden_states)) + return hidden_states + + +class CogVLM2VisionTransformer(nn.Module): + def __init__(self, config): + super().__init__() + self.layers = nn.ModuleList([CogVLM2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +class CogVLM2VisionProjector(nn.Module): + def __init__(self, config, in_features): + super().__init__() + self.linear_proj = nn.Linear(in_features, config.hidden_size, bias=False) + self.norm1 = nn.LayerNorm(config.hidden_size) + self.act1 = nn.GELU() + self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, hidden_states): + hidden_states = self.linear_proj(hidden_states) + hidden_states = self.act1(self.norm1(hidden_states)) + hidden_states = F.silu(self.gate_proj(hidden_states)) * self.dense_h_to_4h(hidden_states) + return self.dense_4h_to_h(hidden_states) + + +class CogVLM2VisionModel(nn.Module): + def __init__(self, config): + super().__init__() + vision_config = config.vision_config + self.patch_embedding = CogVLM2PatchEmbedding(vision_config) + self.transformer = CogVLM2VisionTransformer(vision_config) + self.linear_proj = CogVLM2VisionProjector(config, vision_config.hidden_size) + self.conv = nn.Conv2d( + vision_config.hidden_size, + vision_config.hidden_size, + kernel_size=2, + stride=2, + ) + self.boi = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.eoi = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + + def forward(self, pixel_values): + hidden_states = self.patch_embedding(pixel_values) + hidden_states = self.transformer(hidden_states)[:, 1:] + batch_size, seq_len, hidden_size = hidden_states.shape + grid_size = math.isqrt(seq_len) + if grid_size * grid_size != seq_len: + raise ValueError(f"Vision patch sequence length {seq_len} is not square.") + hidden_states = hidden_states.view(batch_size, grid_size, grid_size, hidden_size).permute(0, 3, 1, 2) + hidden_states = self.conv(hidden_states).flatten(2).transpose(1, 2) + hidden_states = self.linear_proj(hidden_states) + boi = self.boi.expand(batch_size, -1, -1) + eoi = self.eoi.expand(batch_size, -1, -1) + return torch.cat((boi, hidden_states, eoi), dim=1) + + +class CogVLM2LanguageMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + + +class CogVLM2VisionExpertMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.language_mlp = CogVLM2LanguageMLP(config) + + def forward(self, hidden_states): + return self.language_mlp(hidden_states) + + +class CogVLM2RotaryEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + dim = config.hidden_size // config.num_attention_heads + inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq) + + def forward(self, hidden_states, position_ids): + positions = position_ids.to( + device=self.inv_freq.device, + dtype=self.inv_freq.dtype, + ) + freqs = positions.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) + embeddings = torch.cat((freqs, freqs), dim=-1) + return ( + embeddings.cos().to(hidden_states.dtype), + embeddings.sin().to(hidden_states.dtype), + ) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_forward_from_hub("rotary_pos_emb") +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) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class CogVLM2Attention(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_multi_query_heads + self.head_dim = self.hidden_size // self.num_attention_heads + self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + qkv_size = (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim + self.language_expert_query_key_value = nn.Linear(self.hidden_size, qkv_size, bias=False) + self.language_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.rotary_emb = CogVLM2RotaryEmbedding(config) + + def forward( + self, + hidden_states, + position_ids, + attention_mask=None, + past_key_values=None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size, seq_len, _ = hidden_states.shape + qkv = self.language_expert_query_key_value(hidden_states) + query_size = self.num_attention_heads * self.head_dim + kv_size = self.num_key_value_heads * self.head_dim + query, key, value = qkv.split((query_size, kv_size, kv_size), dim=-1) + + query = query.view(batch_size, seq_len, self.num_attention_heads, self.head_dim).transpose(1, 2) + key = key.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value = value.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + cos, sin = self.rotary_emb(hidden_states, position_ids) + query, key = apply_rotary_pos_emb(query, key, cos, sin) + + if past_key_values is not None: + key, value = past_key_values.update(key, value, self.layer_idx) + + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + attn_output, attn_weights = attention_interface( + self, + query, + key, + value, + attention_mask, + scaling=self.scaling, + dropout=0.0, + **kwargs, + ) + attn_output = attn_output.reshape(batch_size, seq_len, -1).contiguous() + return self.language_expert_dense(attn_output), attn_weights + + +class CogVLM2DecoderLayer(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + self.self_attn = CogVLM2Attention(config, layer_idx) + self.mlp = CogVLM2VisionExpertMLP(config) + self.input_layernorm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states, + position_ids, + attention_mask=None, + past_key_values=None, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class CogVLM2PreTrainedModel(PreTrainedModel): + config: CogVLM2Config + base_model_prefix = "model" + supports_gradient_checkpointing = False + _no_split_modules = ["CogVLM2DecoderLayer", "CogVLM2VisionEncoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_sdpa = True + _supports_attention_backend = True + _can_record_outputs = { + "attentions": CogVLM2Attention, + "hidden_states": CogVLM2DecoderLayer, + } + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.\d+\.self_attn\.rotary_emb\.inv_freq"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +LANGUAGE_TOKEN_TYPE = 0 +VISION_TOKEN_TYPE = 1 + + +def build_position_ids(token_type_ids, attention_mask=None): + """Reproduce CogVLM2's compressed position scheme for visual spans.""" + tmp = token_type_ids.clone() + if attention_mask is not None: + tmp[~attention_mask.bool()] = -1 + + is_boundary = torch.zeros_like(tmp, dtype=torch.bool) + is_boundary[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE) + is_boundary[:, 0] |= tmp[:, 0] == VISION_TOKEN_TYPE + is_boundary[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) + is_boundary[:, -1] |= tmp[:, -1] == VISION_TOKEN_TYPE + tmp[is_boundary] = LANGUAGE_TOKEN_TYPE + + position_ids = torch.zeros_like(tmp, dtype=torch.long) + position_ids[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ( + (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE) + ) + return position_ids.cumsum(dim=-1) + + +@auto_docstring +class CogVLM2Model(CogVLM2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [CogVLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.vision = CogVLM2VisionModel(config) + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def get_video_features(self, pixel_values_videos): + return self.vision(pixel_values_videos) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids=None, + token_type_ids=None, + attention_mask=None, + position_ids=None, + past_key_values: Cache | None = None, + inputs_embeds=None, + pixel_values_videos=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ): + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds.") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if token_type_ids is None: + token_type_ids = torch.full( + (inputs_embeds.shape[0], inputs_embeds.shape[1]), + LANGUAGE_TOKEN_TYPE, + dtype=torch.long, + device=inputs_embeds.device, + ) + + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if pixel_values_videos is not None and past_seen_tokens == 0: + video_features = self.get_video_features(pixel_values_videos) + video_features = video_features.reshape(-1, self.config.hidden_size) + vision_mask = token_type_ids[:, -inputs_embeds.shape[1] :] == VISION_TOKEN_TYPE + expected = int(vision_mask.sum().item()) + if expected != video_features.shape[0]: + raise ValueError( + "CogVLM2 visual placeholder count does not match encoded video features: " + f"{expected} placeholders vs {video_features.shape[0]} features." + ) + inputs_embeds = inputs_embeds.clone() + inputs_embeds[vision_mask] = video_features.to(inputs_embeds.device, inputs_embeds.dtype) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + past_seen_tokens = 0 + + if position_ids is None: + if token_type_ids is not None: + full_position_ids = build_position_ids(token_type_ids, attention_mask) + position_ids = full_position_ids[:, -inputs_embeds.shape[1] :] + else: + position_ids = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ).unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + position_ids=position_ids, + attention_mask=causal_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values) + + +@auto_docstring +class CogVLM2ForConditionalGeneration(CogVLM2PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = CogVLM2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, value): + self.lm_head = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids=None, + token_type_ids=None, + attention_mask=None, + position_ids=None, + past_key_values: Cache | None = None, + inputs_embeds=None, + pixel_values_videos=None, + labels=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ): + outputs = self.model( + input_ids=input_ids, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values_videos=pixel_values_videos, + use_cache=use_cache, + **kwargs, + ) + logits = self.lm_head(outputs.last_hidden_state) + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + position_ids=None, + token_type_ids=None, + pixel_values_videos=None, + past_key_values=None, + use_cache=False, + is_first_iteration=False, + **kwargs, + ): + if position_ids is None and token_type_ids is not None: + position_ids = build_position_ids(token_type_ids, attention_mask) + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + use_cache=use_cache, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values_videos=pixel_values_videos, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if not is_first_iteration and use_cache: + model_inputs["pixel_values_videos"] = None + + return model_inputs + + +__all__ = ["CogVLM2VisionModel", "CogVLM2PreTrainedModel", "CogVLM2Model", "CogVLM2ForConditionalGeneration"] diff --git a/src/transformers/models/cogvlm2/modular_cogvlm2.py b/src/transformers/models/cogvlm2/modular_cogvlm2.py new file mode 100644 index 000000000000..c4afcb94d552 --- /dev/null +++ b/src/transformers/models/cogvlm2/modular_cogvlm2.py @@ -0,0 +1,655 @@ +# Copyright 2026 The HuggingFace 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. + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig +from ...generation import GenerationMixin +from ...masking_utils import create_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..llama.modeling_llama import apply_rotary_pos_emb, eager_attention_forward + + +LANGUAGE_TOKEN_TYPE = 0 +VISION_TOKEN_TYPE = 1 + + +class CogVLM2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class CogVLM2VisionConfig(PreTrainedConfig): + model_type = "cogvlm2_vision" + base_config_key = "vision_config" + + def __init__( + self, + hidden_size=1792, + image_size=224, + in_channels=3, + intermediate_size=15360, + layer_norm_eps=1e-6, + num_heads=16, + num_hidden_layers=63, + num_positions=257, + patch_size=14, + dropout_prob=0.0, + hidden_act="gelu", + **kwargs, + ): + self.hidden_size = hidden_size + self.image_size = image_size + self.in_channels = in_channels + self.intermediate_size = intermediate_size + self.layer_norm_eps = layer_norm_eps + self.num_heads = num_heads + self.num_hidden_layers = num_hidden_layers + self.num_positions = num_positions + self.patch_size = patch_size + self.dropout_prob = dropout_prob + self.hidden_act = hidden_act + super().__init__(**kwargs) + + +class CogVLM2Config(PreTrainedConfig): + model_type = "cogvlm2" + sub_configs = {"vision_config": CogVLM2VisionConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=128256, + hidden_size=4096, + intermediate_size=14336, + num_hidden_layers=32, + num_attention_heads=32, + num_multi_query_heads=8, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-5, + template_version="base", + vision_config=None, + rope_theta=500000.0, + pad_token_id=128002, + bos_token_id=128000, + eos_token_id=128001, + tie_word_embeddings=False, + use_cache=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_multi_query_heads = num_multi_query_heads + self.num_key_value_heads = num_multi_query_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.template_version = template_version + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} + + if vision_config is None: + self.vision_config = CogVLM2VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = CogVLM2VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +def build_position_ids(token_type_ids, attention_mask=None): + """Reproduce CogVLM2's compressed position scheme for visual spans.""" + tmp = token_type_ids.clone() + if attention_mask is not None: + tmp[~attention_mask.bool()] = -1 + + is_boundary = torch.zeros_like(tmp, dtype=torch.bool) + is_boundary[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE) + is_boundary[:, 0] |= tmp[:, 0] == VISION_TOKEN_TYPE + is_boundary[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) + is_boundary[:, -1] |= tmp[:, -1] == VISION_TOKEN_TYPE + tmp[is_boundary] = LANGUAGE_TOKEN_TYPE + + position_ids = torch.zeros_like(tmp, dtype=torch.long) + position_ids[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ( + (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE) + ) + return position_ids.cumsum(dim=-1) + + +class CogVLM2PatchEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + self.proj = nn.Conv2d( + config.in_channels, + config.hidden_size, + kernel_size=config.patch_size, + stride=config.patch_size, + ) + self.cls_embedding = nn.Parameter(torch.zeros(1, config.hidden_size)) + self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size) + + def forward(self, pixel_values): + hidden_states = self.proj(pixel_values).flatten(2).transpose(1, 2) + cls_token = self.cls_embedding.expand(hidden_states.shape[0], -1, -1) + hidden_states = torch.cat((cls_token, hidden_states), dim=1) + return hidden_states + self.position_embedding.weight.unsqueeze(0) + + +class CogVLM2VisionAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.hidden_size // config.num_heads + self.query_key_value = nn.Linear(config.hidden_size, config.hidden_size * 3) + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.output_dropout = nn.Dropout(config.dropout_prob) + + def forward(self, hidden_states): + batch_size, seq_len, _ = hidden_states.shape + qkv = self.query_key_value(hidden_states) + qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim) + query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) + attn_output = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1) + return self.output_dropout(self.dense(attn_output)) + + +class CogVLM2VisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states): + return self.fc2(self.activation_fn(self.fc1(hidden_states))) + + +class CogVLM2VisionEncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attention = CogVLM2VisionAttention(config) + self.mlp = CogVLM2VisionMLP(config) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states): + residual = hidden_states + hidden_states = residual + self.input_layernorm(self.attention(hidden_states)) + residual = hidden_states + hidden_states = residual + self.post_attention_layernorm(self.mlp(hidden_states)) + return hidden_states + + +class CogVLM2VisionTransformer(nn.Module): + def __init__(self, config): + super().__init__() + self.layers = nn.ModuleList([CogVLM2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward(self, hidden_states): + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states +class CogVLM2VisionProjector(nn.Module): + def __init__(self, config, in_features): + super().__init__() + self.linear_proj = nn.Linear(in_features, config.hidden_size, bias=False) + self.norm1 = nn.LayerNorm(config.hidden_size) + self.act1 = nn.GELU() + self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, hidden_states): + hidden_states = self.linear_proj(hidden_states) + hidden_states = self.act1(self.norm1(hidden_states)) + hidden_states = F.silu(self.gate_proj(hidden_states)) * self.dense_h_to_4h(hidden_states) + return self.dense_4h_to_h(hidden_states) + + +class CogVLM2VisionModel(nn.Module): + def __init__(self, config): + super().__init__() + vision_config = config.vision_config + self.patch_embedding = CogVLM2PatchEmbedding(vision_config) + self.transformer = CogVLM2VisionTransformer(vision_config) + self.linear_proj = CogVLM2VisionProjector(config, vision_config.hidden_size) + self.conv = nn.Conv2d( + vision_config.hidden_size, + vision_config.hidden_size, + kernel_size=2, + stride=2, + ) + self.boi = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + self.eoi = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) + + def forward(self, pixel_values): + hidden_states = self.patch_embedding(pixel_values) + hidden_states = self.transformer(hidden_states)[:, 1:] + batch_size, seq_len, hidden_size = hidden_states.shape + grid_size = math.isqrt(seq_len) + if grid_size * grid_size != seq_len: + raise ValueError(f"Vision patch sequence length {seq_len} is not square.") + hidden_states = hidden_states.view(batch_size, grid_size, grid_size, hidden_size).permute(0, 3, 1, 2) + hidden_states = self.conv(hidden_states).flatten(2).transpose(1, 2) + hidden_states = self.linear_proj(hidden_states) + boi = self.boi.expand(batch_size, -1, -1) + eoi = self.eoi.expand(batch_size, -1, -1) + return torch.cat((boi, hidden_states, eoi), dim=1) + + +class CogVLM2LanguageMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + + +class CogVLM2VisionExpertMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.language_mlp = CogVLM2LanguageMLP(config) + + def forward(self, hidden_states): + return self.language_mlp(hidden_states) + + +class CogVLM2RotaryEmbedding(nn.Module): + def __init__(self, config): + super().__init__() + dim = config.hidden_size // config.num_attention_heads + inv_freq = 1.0 / ( + config.rope_theta + ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + self.register_buffer("inv_freq", inv_freq) + + def forward(self, hidden_states, position_ids): + positions = position_ids.to( + device=self.inv_freq.device, + dtype=self.inv_freq.dtype, + ) + freqs = positions.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) + embeddings = torch.cat((freqs, freqs), dim=-1) + return ( + embeddings.cos().to(hidden_states.dtype), + embeddings.sin().to(hidden_states.dtype), + ) + + +class CogVLM2Attention(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_multi_query_heads + self.head_dim = self.hidden_size // self.num_attention_heads + self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads + self.scaling = self.head_dim**-0.5 + qkv_size = (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim + self.language_expert_query_key_value = nn.Linear(self.hidden_size, qkv_size, bias=False) + self.language_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + self.rotary_emb = CogVLM2RotaryEmbedding(config) + + def forward( + self, + hidden_states, + position_ids, + attention_mask=None, + past_key_values=None, + **kwargs: Unpack[TransformersKwargs], + ): + batch_size, seq_len, _ = hidden_states.shape + qkv = self.language_expert_query_key_value(hidden_states) + query_size = self.num_attention_heads * self.head_dim + kv_size = self.num_key_value_heads * self.head_dim + query, key, value = qkv.split((query_size, kv_size, kv_size), dim=-1) + + query = query.view(batch_size, seq_len, self.num_attention_heads, self.head_dim).transpose(1, 2) + key = key.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value = value.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + cos, sin = self.rotary_emb(hidden_states, position_ids) + query, key = apply_rotary_pos_emb(query, key, cos, sin) + + if past_key_values is not None: + key, value = past_key_values.update(key, value, self.layer_idx) + + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + attn_output, attn_weights = attention_interface( + self, + query, + key, + value, + attention_mask, + scaling=self.scaling, + dropout=0.0, + **kwargs, + ) + attn_output = attn_output.reshape(batch_size, seq_len, -1).contiguous() + return self.language_expert_dense(attn_output), attn_weights + + +class CogVLM2DecoderLayer(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + self.self_attn = CogVLM2Attention(config, layer_idx) + self.mlp = CogVLM2VisionExpertMLP(config) + self.input_layernorm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states, + position_ids, + attention_mask=None, + past_key_values=None, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class CogVLM2PreTrainedModel(PreTrainedModel): + config: CogVLM2Config + base_model_prefix = "model" + supports_gradient_checkpointing = False + _no_split_modules = ["CogVLM2DecoderLayer", "CogVLM2VisionEncoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_sdpa = True + _supports_attention_backend = True + _can_record_outputs = { + "attentions": CogVLM2Attention, + "hidden_states": CogVLM2DecoderLayer, + } + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.\d+\.self_attn\.rotary_emb\.inv_freq"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +@auto_docstring +class CogVLM2Model(CogVLM2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [CogVLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = CogVLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.vision = CogVLM2VisionModel(config) + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def get_video_features(self, pixel_values_videos): + return self.vision(pixel_values_videos) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids=None, + token_type_ids=None, + attention_mask=None, + position_ids=None, + past_key_values: Cache | None = None, + inputs_embeds=None, + pixel_values_videos=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ): + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds.") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if token_type_ids is None: + token_type_ids = torch.full( + (inputs_embeds.shape[0], inputs_embeds.shape[1]), + LANGUAGE_TOKEN_TYPE, + dtype=torch.long, + device=inputs_embeds.device, + ) + + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if pixel_values_videos is not None and past_seen_tokens == 0: + video_features = self.get_video_features(pixel_values_videos) + video_features = video_features.reshape(-1, self.config.hidden_size) + vision_mask = token_type_ids[:, -inputs_embeds.shape[1] :] == VISION_TOKEN_TYPE + expected = int(vision_mask.sum().item()) + if expected != video_features.shape[0]: + raise ValueError( + "CogVLM2 visual placeholder count does not match encoded video features: " + f"{expected} placeholders vs {video_features.shape[0]} features." + ) + inputs_embeds = inputs_embeds.clone() + inputs_embeds[vision_mask] = video_features.to(inputs_embeds.device, inputs_embeds.dtype) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + past_seen_tokens = 0 + + if position_ids is None: + if token_type_ids is not None: + full_position_ids = build_position_ids(token_type_ids, attention_mask) + position_ids = full_position_ids[:, -inputs_embeds.shape[1] :] + else: + position_ids = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ).unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + position_ids=position_ids, + attention_mask=causal_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values) + + +@auto_docstring +class CogVLM2ForConditionalGeneration(CogVLM2PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = CogVLM2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, value): + self.lm_head = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids=None, + token_type_ids=None, + attention_mask=None, + position_ids=None, + past_key_values: Cache | None = None, + inputs_embeds=None, + pixel_values_videos=None, + labels=None, + use_cache=None, + **kwargs: Unpack[TransformersKwargs], + ): + outputs = self.model( + input_ids=input_ids, + token_type_ids=token_type_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values_videos=pixel_values_videos, + use_cache=use_cache, + **kwargs, + ) + logits = self.lm_head(outputs.last_hidden_state) + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids=None, + inputs_embeds=None, + attention_mask=None, + position_ids=None, + token_type_ids=None, + pixel_values_videos=None, + past_key_values=None, + use_cache=False, + is_first_iteration=False, + **kwargs, + ): + if position_ids is None and token_type_ids is not None: + position_ids = build_position_ids(token_type_ids, attention_mask) + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + use_cache=use_cache, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values_videos=pixel_values_videos, + is_first_iteration=is_first_iteration, + **kwargs, + ) + + if not is_first_iteration and use_cache: + model_inputs["pixel_values_videos"] = None + + return model_inputs + + +__all__ = [ + "CogVLM2VisionConfig", + "CogVLM2Config", + "CogVLM2VisionModel", + "CogVLM2PreTrainedModel", + "CogVLM2Model", + "CogVLM2ForConditionalGeneration", +] diff --git a/src/transformers/models/cogvlm2/processing_cogvlm2.py b/src/transformers/models/cogvlm2/processing_cogvlm2.py new file mode 100644 index 000000000000..845c295760e0 --- /dev/null +++ b/src/transformers/models/cogvlm2/processing_cogvlm2.py @@ -0,0 +1,236 @@ +# Copyright 2026 The HuggingFace 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 ...feature_extraction_utils import BatchFeature +from ...processing_utils import ProcessingKwargs, ProcessorMixin +from ...utils import auto_docstring +from ...video_utils import VideoInput + + +LANGUAGE_TOKEN_TYPE = 0 +VISION_TOKEN_TYPE = 1 + + +class CogVLM2ProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": {"padding": False}, + "videos_kwargs": {"return_metadata": False}, + } + + +@auto_docstring +class CogVLM2Processor(ProcessorMixin): + valid_processor_kwargs = CogVLM2ProcessorKwargs + + def __init__( + self, + video_processor=None, + tokenizer=None, + chat_template=None, + template_version="base", + num_visual_tokens_per_frame=66, + ): + r""" + template_version (`str`, *optional*, defaults to `"base"`): + Default CogVLM2 prompt template. Supported values are `"base"`, `"chat"`, and `"vqa"`. + num_visual_tokens_per_frame (`int`, *optional*, defaults to 66): + Number of visual placeholder tokens per video frame. The released CogVLM2-Video + checkpoint produces 64 spatial tokens plus BOI and EOI tokens. + """ + self.template_version = template_version + self.num_visual_tokens_per_frame = num_visual_tokens_per_frame + self.template_version = template_version + super().__init__(video_processor, tokenizer, chat_template=chat_template) + + @staticmethod + def build_prompt(query, history=None, template_version="chat"): + history = history or [] + if template_version == "base": + return query + if template_version == "vqa": + answer_format = "Short answer:" + elif template_version == "chat": + answer_format = "Answer:" + else: + raise ValueError(f"Unknown template version: {template_version}") + + prompt = "" + for old_query, response in history: + prompt += f"Question: {old_query} {answer_format} {response}\n" + return prompt + f"Question: {query} {answer_format}" + + def __call__( + self, + text=None, + videos: VideoInput | None = None, + *, + history=None, + template_version=None, + add_time_indices=None, + answer=None, + padding=False, + truncation=False, + max_length=None, + return_tensors=None, + **videos_kwargs, + ): + if text is None and videos is None: + raise ValueError("You have to specify at least one of text or videos.") + + template_version = template_version or self.template_version + + if text is None: + text = "" + if isinstance(text, str): + text = [text] + elif not isinstance(text, list) or not all(isinstance(item, str) for item in text): + raise ValueError("text must be a string or a list of strings.") + + template_version = template_version or self.template_version + if add_time_indices is None: + add_time_indices = template_version == "chat" + + if history is not None: + if len(text) != 1: + raise ValueError("history is only supported for a single text input.") + text = [self.build_prompt(text[0], history=history, template_version=template_version)] + elif template_version != "base": + text = [self.build_prompt(item, template_version=template_version) for item in text] + + if answer is None: + answers = [None] * len(text) + elif isinstance(answer, str): + if len(text) != 1: + raise ValueError("A single answer string can only be used with one text sample.") + answers = [answer] + elif isinstance(answer, list) and len(answer) == len(text) and all(isinstance(item, str) for item in answer): + answers = answer + else: + raise ValueError("answer must be a string or a list of strings matching the text batch size.") + + video_inputs = {} + if videos is not None: + video_inputs = self.video_processor(videos, return_tensors=return_tensors, **videos_kwargs) + frame_counts = video_inputs.pop("video_frame_counts") + frame_counts = frame_counts.tolist() if hasattr(frame_counts, "tolist") else list(frame_counts) + if len(frame_counts) != len(text): + raise ValueError( + "CogVLM2 expects one video per text sample: " + f"got {len(frame_counts)} videos and {len(text)} text samples." + ) + else: + frame_counts = [0] * len(text) + pad_token_id = self.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = 128002 + bos_token_id = self.tokenizer.bos_token_id + if bos_token_id is None: + raise ValueError("CogVLM2 requires the tokenizer to define bos_token_id.") + + batch_input_ids = [] + batch_token_type_ids = [] + batch_labels = [] + visual_tokens_per_frame = self.num_visual_tokens_per_frame + + for sample_text, num_frames, sample_answer in zip(text, frame_counts, answers): + input_ids = [bos_token_id] + token_type_ids = [LANGUAGE_TOKEN_TYPE] + + for frame_idx in range(num_frames): + input_ids.extend([pad_token_id] * visual_tokens_per_frame) + token_type_ids.extend([VISION_TOKEN_TYPE] * visual_tokens_per_frame) + if add_time_indices: + time_ids = self.tokenizer.encode(str(frame_idx), add_special_tokens=False) + input_ids.extend(time_ids) + token_type_ids.extend([LANGUAGE_TOKEN_TYPE] * len(time_ids)) + + multimodal_prefix_length = len(input_ids) + text_ids = self.tokenizer.encode(sample_text, add_special_tokens=False) + input_ids.extend(text_ids) + token_type_ids.extend([LANGUAGE_TOKEN_TYPE] * len(text_ids)) + + if sample_answer is not None: + answer_ids = self.tokenizer.encode(sample_answer, add_special_tokens=False) + [self.tokenizer.eos_token_id] + labels = [-100] * len(input_ids) + answer_ids + input_ids.extend(answer_ids) + token_type_ids.extend([LANGUAGE_TOKEN_TYPE] * len(answer_ids)) + else: + labels = None + + if truncation: + effective_max_length = max_length if max_length is not None else self.tokenizer.model_max_length + if len(input_ids) > effective_max_length: + if effective_max_length < multimodal_prefix_length: + raise ValueError( + "Truncation would remove CogVLM2 visual placeholder tokens. " + "Increase max_length or reduce the number of video frames." + ) + if self.tokenizer.truncation_side == "left" and num_frames: + raise ValueError( + "Left truncation is not supported for CogVLM2 video inputs because it would " + "remove visual placeholders from the beginning of the sequence." + ) + if self.tokenizer.truncation_side == "left": + input_ids = input_ids[-effective_max_length:] + token_type_ids = token_type_ids[-effective_max_length:] + if labels is not None: + labels = labels[-effective_max_length:] + else: + input_ids = input_ids[:effective_max_length] + token_type_ids = token_type_ids[:effective_max_length] + if labels is not None: + labels = labels[:effective_max_length] + + batch_input_ids.append(input_ids) + batch_token_type_ids.append(token_type_ids) + batch_labels.append(labels) + encoded = self.tokenizer.pad( + {"input_ids": batch_input_ids, "token_type_ids": batch_token_type_ids}, + padding=padding, + max_length=max_length if padding == "max_length" else None, + return_attention_mask=True, + return_tensors=return_tensors, + ) + + data = dict(encoded) + if any(labels is not None for labels in batch_labels): + if not all(labels is not None for labels in batch_labels): + raise ValueError("Either all samples in a batch must provide answer labels or none of them.") + + if hasattr(encoded["input_ids"], "shape"): + padded_length = encoded["input_ids"].shape[-1] + else: + padded_length = max(len(ids) for ids in encoded["input_ids"]) + + padded_labels = [] + for labels in batch_labels: + difference = padded_length - len(labels) + if self.tokenizer.padding_side == "left": + labels = [-100] * difference + labels + else: + labels = labels + [-100] * difference + padded_labels.append(labels) + data["labels"] = padded_labels + + data.update(video_inputs) + return BatchFeature(data=data, tensor_type=return_tensors) + + @property + def model_input_names(self): + tokenizer_input_names = getattr(self.tokenizer, "model_input_names", []) + video_input_names = getattr(self.video_processor, "model_input_names", []) + return list(dict.fromkeys(tokenizer_input_names + ["token_type_ids"] + video_input_names)) + + +__all__ = ["CogVLM2Processor"] diff --git a/src/transformers/models/cogvlm2/video_processing_cogvlm2.py b/src/transformers/models/cogvlm2/video_processing_cogvlm2.py new file mode 100644 index 000000000000..dcf43d8790aa --- /dev/null +++ b/src/transformers/models/cogvlm2/video_processing_cogvlm2.py @@ -0,0 +1,95 @@ +# Copyright 2026 The HuggingFace 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. + +import torch +import torchvision.transforms.v2.functional as tvF + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict +from ...utils import TensorType +from ...video_processing_utils import BaseVideoProcessor +from ...video_utils import group_videos_by_shape, reorder_videos + + +class CogVLM2VideoProcessor(BaseVideoProcessor): + resample = PILImageResampling.BICUBIC + image_mean = OPENAI_CLIP_MEAN + image_std = OPENAI_CLIP_STD + size = {"shortest_edge": 224} + crop_size = {"height": 224, "width": 224} + do_resize = True + do_center_crop = True + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_convert_rgb = True + do_sample_frames = False + return_metadata = False + model_input_names = ["pixel_values_videos"] + + def _preprocess( + self, + videos: list["torch.Tensor"], + do_convert_rgb: bool, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + for shape, stacked_videos in grouped_videos.items(): + if do_convert_rgb: + stacked_videos = self.convert_to_rgb(stacked_videos) + if do_resize: + stacked_videos = self.resize(stacked_videos, size=size, resample=resample) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + for shape, stacked_videos in grouped_videos.items(): + if do_center_crop: + stacked_videos = self.center_crop(stacked_videos, crop_size) + stacked_videos = self.rescale_and_normalize( + stacked_videos, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + processed_videos_grouped[shape] = stacked_videos + + processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index) + frame_counts = torch.tensor([video.shape[0] for video in processed_videos], dtype=torch.long) + pixel_values_videos = torch.cat(processed_videos, dim=0) + return BatchFeature( + data={ + "pixel_values_videos": pixel_values_videos, + "video_frame_counts": frame_counts, + }, + tensor_type=return_tensors, + ) + + +__all__ = ["CogVLM2VideoProcessor"] diff --git a/tests/models/cogvlm2/__init__.py b/tests/models/cogvlm2/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/cogvlm2/test_modeling_cogvlm2.py b/tests/models/cogvlm2/test_modeling_cogvlm2.py new file mode 100644 index 000000000000..49f0d31248df --- /dev/null +++ b/tests/models/cogvlm2/test_modeling_cogvlm2.py @@ -0,0 +1,128 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. + +import unittest + +import torch + +from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + CogVLM2Config, + CogVLM2ForConditionalGeneration, +) +from transformers.models.cogvlm2.modeling_cogvlm2 import build_position_ids +from transformers.testing_utils import require_torch + + +@require_torch +class CogVLM2ModelTest(unittest.TestCase): + def get_config(self): + return CogVLM2Config( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_multi_query_heads=2, + max_position_embeddings=128, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + vision_config={ + "hidden_size": 32, + "image_size": 28, + "in_channels": 3, + "intermediate_size": 64, + "num_heads": 4, + "num_hidden_layers": 2, + "num_positions": 17, + "patch_size": 7, + }, + ) + + def get_inputs(self): + input_ids = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 5, 6, 7]]) + token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 0, 0, 0]]) + return { + "input_ids": input_ids, + "token_type_ids": token_type_ids, + "attention_mask": torch.ones_like(input_ids), + "pixel_values_videos": torch.randn(1, 3, 28, 28), + } + + def test_position_ids_compress_visual_span(self): + token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 0, 0, 0]]) + position_ids = build_position_ids(token_type_ids) + self.assertEqual(position_ids.tolist(), [[0, 1, 2, 2, 2, 2, 3, 4, 5, 6]]) + + def test_forward_and_cache(self): + model = CogVLM2ForConditionalGeneration(self.get_config()).eval() + inputs = self.get_inputs() + with torch.no_grad(): + outputs = model(**inputs, use_cache=True) + self.assertEqual(outputs.logits.shape, (1, 10, 128)) + self.assertEqual(outputs.past_key_values.get_seq_length(), 10) + self.assertTrue(torch.isfinite(outputs.logits).all()) + + def test_visual_placeholder_mismatch_raises(self): + model = CogVLM2ForConditionalGeneration(self.get_config()).eval() + inputs = self.get_inputs() + inputs["token_type_ids"][0, 6] = 0 + with self.assertRaisesRegex(ValueError, "visual placeholder count"): + model(**inputs) + + def test_hidden_states_and_eager_attentions(self): + config = self.get_config() + config._attn_implementation = "eager" + model = CogVLM2ForConditionalGeneration(config).eval() + with torch.no_grad(): + outputs = model( + **self.get_inputs(), + output_hidden_states=True, + output_attentions=True, + use_cache=False, + ) + self.assertEqual(len(outputs.hidden_states), config.num_hidden_layers + 1) + self.assertEqual(len(outputs.attentions), config.num_hidden_layers) + self.assertEqual(outputs.attentions[0].shape, (1, 4, 10, 10)) + + def test_generate_with_video_and_cache(self): + model = CogVLM2ForConditionalGeneration(self.get_config()).eval() + inputs = self.get_inputs() + with torch.no_grad(): + outputs = model.generate(**inputs, max_new_tokens=2, do_sample=False) + self.assertEqual(outputs.shape, (1, 12)) + + def test_auto_classes(self): + config = AutoConfig.for_model( + "cogvlm2", + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_multi_query_heads=2, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + vision_config={ + "hidden_size": 32, + "image_size": 28, + "in_channels": 3, + "intermediate_size": 64, + "num_heads": 4, + "num_hidden_layers": 1, + "num_positions": 17, + "patch_size": 7, + }, + ) + self.assertEqual(type(AutoModel.from_config(config)).__name__, "CogVLM2Model") + self.assertEqual( + type(AutoModelForCausalLM.from_config(config)).__name__, + "CogVLM2ForConditionalGeneration", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/cogvlm2/test_processing_cogvlm2.py b/tests/models/cogvlm2/test_processing_cogvlm2.py new file mode 100644 index 000000000000..0e8d9d168472 --- /dev/null +++ b/tests/models/cogvlm2/test_processing_cogvlm2.py @@ -0,0 +1,110 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. + +import unittest + +import torch + +from transformers import CogVLM2Processor, CogVLM2VideoProcessor, PreTrainedTokenizer +from transformers.testing_utils import require_torch, require_torchvision + + +class TinyCogVLM2Tokenizer(PreTrainedTokenizer): + model_input_names = ["input_ids", "attention_mask"] + + def __init__(self): + chars = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 :?.,!\\n") + self.vocab = {"": 0, "": 1, "": 2, "": 3} + for char in chars: + if char not in self.vocab: + self.vocab[char] = len(self.vocab) + self.inverse_vocab = {value: key for key, value in self.vocab.items()} + super().__init__(pad_token="", bos_token="", eos_token="", unk_token="") + + @property + def vocab_size(self): + return len(self.vocab) + + def get_vocab(self): + return dict(self.vocab) + + def _tokenize(self, text): + return list(text) + + def _convert_token_to_id(self, token): + return self.vocab.get(token, self.vocab[""]) + + def _convert_id_to_token(self, index): + return self.inverse_vocab.get(index, "") + + def convert_tokens_to_string(self, tokens): + return "".join(tokens) + + def save_vocabulary(self, save_directory, filename_prefix=None): + return () + + +@require_torch +@require_torchvision +class CogVLM2ProcessorTest(unittest.TestCase): + def setUp(self): + self.tokenizer = TinyCogVLM2Tokenizer() + self.video_processor = CogVLM2VideoProcessor() + self.processor = CogVLM2Processor( + video_processor=self.video_processor, + tokenizer=self.tokenizer, + ) + + def test_builds_66_visual_tokens_per_frame(self): + video = torch.randint(0, 256, (2, 3, 240, 320), dtype=torch.uint8) + outputs = self.processor(text="What?", videos=video, return_tensors="pt") + + self.assertEqual(int((outputs.token_type_ids == 1).sum()), 132) + self.assertEqual(outputs.pixel_values_videos.shape, (2, 3, 224, 224)) + + def test_base_prompt_matches_config_default(self): + video = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.uint8) + outputs = self.processor(text="What?", videos=video, template_version="chat", return_tensors="pt") + + text_start = 1 + 66 + decoded = self.tokenizer.decode(outputs.input_ids[0, text_start:], skip_special_tokens=True) + self.assertEqual(decoded, "What?") + + def test_chat_prompt_matches_reference_format_and_adds_time_index(self): + video = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.uint8) + outputs = self.processor( + text="What?", + videos=video, + template_version="chat", + return_tensors="pt", + ) + + language_ids = outputs.input_ids[0][outputs.token_type_ids[0] == 0] + decoded = self.tokenizer.decode(language_ids, skip_special_tokens=True) + self.assertEqual(decoded, "0Question: What? Answer:") + + def test_default_template_matches_checkpoint_base_mode(self): + video = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.uint8) + outputs = self.processor(text="What?", videos=video, return_tensors="pt") + + text_start = 1 + 66 + decoded = self.tokenizer.decode(outputs.input_ids[0, text_start:], skip_special_tokens=True) + self.assertEqual(decoded, "What?") + + def test_answer_labels_only_supervise_answer(self): + video = torch.randint(0, 256, (1, 3, 224, 224), dtype=torch.uint8) + outputs = self.processor(text="What?", videos=video, answer="OK", return_tensors="pt") + + self.assertTrue((outputs.labels[0, :-3] == -100).all()) + self.assertEqual(int((outputs.labels != -100).sum()), 3) + + def test_default_does_not_insert_time_indices(self): + video = torch.randint(0, 256, (2, 3, 224, 224), dtype=torch.uint8) + outputs = self.processor(text="X", videos=video, return_tensors="pt") + expected_prefix_length = 1 + 2 * 66 + + first_text_id = self.tokenizer.encode("X", add_special_tokens=False)[0] + self.assertEqual(outputs.input_ids[0, expected_prefix_length].item(), first_text_id) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/cogvlm2/test_video_processing_cogvlm2.py b/tests/models/cogvlm2/test_video_processing_cogvlm2.py new file mode 100644 index 000000000000..8df9209f2dc2 --- /dev/null +++ b/tests/models/cogvlm2/test_video_processing_cogvlm2.py @@ -0,0 +1,38 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. + +import unittest + +import torch + +from transformers import CogVLM2VideoProcessor +from transformers.testing_utils import require_torch, require_torchvision + + +@require_torch +@require_torchvision +class CogVLM2VideoProcessorTest(unittest.TestCase): + def setUp(self): + self.processor = CogVLM2VideoProcessor() + + def test_video_preprocessing_shape_and_dtype(self): + video = torch.randint(0, 256, (2, 3, 240, 320), dtype=torch.uint8) + outputs = self.processor(video, return_tensors="pt") + + self.assertEqual(outputs.pixel_values_videos.shape, (2, 3, 224, 224)) + self.assertEqual(outputs.pixel_values_videos.dtype, torch.float32) + self.assertEqual(outputs.video_frame_counts.tolist(), [2]) + self.assertTrue(torch.isfinite(outputs.pixel_values_videos).all()) + + def test_multiple_videos_preserve_frame_counts(self): + videos = [ + torch.randint(0, 256, (2, 3, 240, 320), dtype=torch.uint8), + torch.randint(0, 256, (3, 3, 320, 240), dtype=torch.uint8), + ] + outputs = self.processor(videos, return_tensors="pt") + + self.assertEqual(outputs.pixel_values_videos.shape, (5, 3, 224, 224)) + self.assertEqual(outputs.video_frame_counts.tolist(), [2, 3]) + + +if __name__ == "__main__": + unittest.main()