# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/hunyuan_vl/modular_hunyuan_vl.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_hunyuan_vl.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright (C) 2026 THL A29 Limited, a Tencent company and the 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. import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Optional import torch from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.deprecation import deprecate_kwarg from ...utils.generic import ( accepts_precomputed_kwargs, is_flash_attention_requested, maybe_autocast, merge_with_config_defaults, ) from ...utils.output_capturing import capture_outputs from ...vision_utils import get_vision_cu_seqlens from .configuration_hunyuan_vl import HunYuanVLConfig, HunYuanVLTextConfig, HunYuanVLVisionConfig @dataclass class HunYuanVLModelOutputWithPast(BaseModelOutputWithPast): r""" image_hidden_states (`torch.FloatTensor`, *optional*): Last image features produced by the vision tower and scattered into the language-model token stream. """ image_hidden_states: torch.FloatTensor | None = None @use_kernel_forward_from_hub("RMSNorm") class HunYuanVLRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ HunYuanVLRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 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) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class HunYuanVLRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` def __init__(self, config: HunYuanVLTextConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] # Diff from Llama - DynamicNTKAlphaRotary if self.rope_type == "dynamic" and self.config.rope_parameters.get("alpha"): self.dim = config.head_dim base = self.config.rope_parameters["rope_theta"] * self.config.rope_parameters["alpha"] ** ( self.config.head_dim / (self.config.head_dim - 2) ) inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.config.head_dim)) self.attention_scaling = 1.0 else: rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) rope_parameters = getattr(config, "rope_parameters", None) or {} self.mrope_section = rope_parameters.get("mrope_section") @staticmethod def compute_default_rope_parameters( config: HunYuanVLConfig | None = None, device: Optional["torch.device"] = None, seq_len: int | None = None, ) -> tuple["torch.Tensor", float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. device (`torch.device`): The device to use for initialization of the inverse frequencies. seq_len (`int`, *optional*): The current sequence length. Unused for this type of RoPE. 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"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads 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.int64).to(device=device, dtype=torch.float) / dim) ) return inv_freq, attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): # In contrast to other models, model has different position ids for the grids # So we expand the inv_freq to shape (3, ...) inv_freq_expanded = ( self.inv_freq[None, None, :, None].float().expand(len(self.mrope_section), position_ids.shape[1], -1, 1) ) position_ids_expanded = position_ids[:, :, None, :].float() # shape (mrope_section, bs, 1, positions) device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with maybe_autocast(device_type=device_type, enabled=False): # Force float32 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class HunYuanVLVisionMLP(nn.Module): def __init__(self, config): super().__init__() self.config = config 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: torch.Tensor) -> torch.Tensor: hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) return hidden_states class HunYuanVLVisionPatchEmbed(nn.Module): def __init__(self, config: HunYuanVLVisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.patch_size = config.patch_size self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, bias=True, ) self.max_num_patches = (config.max_image_size // self.patch_size) ** 2 self.num_positions = self.max_num_patches + 1 self.position_edge = config.max_image_size // self.patch_size # The first token is the cls token; the remaining tokens form the learnable patch positional grid. self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: """Interpolate the learned patch positional grid to the current image grid size.""" patch_pos_embed = self.position_embedding.weight[1:, :].reshape( 1, self.position_edge, self.position_edge, self.embed_dim ) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2).float() patch_pos_embed = nn.functional.interpolate( patch_pos_embed, size=(height, width), mode=self.config.interpolate_mode, align_corners=False, ) return patch_pos_embed.permute(0, 2, 3, 1).reshape(1, -1, self.embed_dim).to(embeddings.dtype) def forward(self, pixel_values: torch.Tensor, grid_thw: list[list[int]]) -> torch.Tensor: num_patches, _ = pixel_values.shape pixel_values = pixel_values.reshape(1, num_patches, self.config.num_channels, self.patch_size, self.patch_size) batch_size, sequence_len, channel, height, width = pixel_values.shape target_dtype = self.patch_embedding.weight.dtype pixel_values = pixel_values.reshape(batch_size * sequence_len, channel, height, width) patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) embeddings = patch_embeds.flatten(-2).squeeze(-1) embeddings = embeddings.reshape(batch_size, sequence_len, -1).squeeze(0) start = 0 image_embeddings_list = [] for t, h, w in grid_thw: end = start + t * h * w image_embeddings = embeddings[start:end, :] position_embedding = self.interpolate_pos_encoding(image_embeddings, h, w).squeeze(0).repeat(t, 1) image_embeddings_list.append(image_embeddings + position_embedding) start = end return torch.concat(image_embeddings_list, dim=0).unsqueeze(0) class HunYuanVLVisionPatchMerger(nn.Module): def __init__(self, config: HunYuanVLVisionConfig): super().__init__() self.config = config self.spatial_merge_size = config.spatial_merge_size self.proj_conv = nn.Conv2d( config.hidden_size, config.hidden_size * 2, kernel_size=config.spatial_merge_size, stride=config.spatial_merge_size, ) self.proj_act = ACT2FN[config.hidden_act] self.proj_out = nn.Conv2d(config.hidden_size * 2, config.hidden_size * 4, kernel_size=1) self.mlp = nn.Linear(config.hidden_size * 4, config.text_hidden_size) self.image_newline = nn.Parameter(torch.empty(config.hidden_size * 4)) self.image_begin = nn.Parameter(torch.empty(config.text_hidden_size)) self.image_end = nn.Parameter(torch.empty(config.text_hidden_size)) self.image_sep = nn.Parameter(torch.empty(config.text_hidden_size)) self.before_rms = HunYuanVLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.after_rms = HunYuanVLRMSNorm(config.text_hidden_size, eps=config.rms_norm_eps) def forward(self, hidden_states: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: hidden_states = self.before_rms(hidden_states) dtype = hidden_states.dtype hidden_states = hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, *size) hidden_states = self.proj_conv(hidden_states) hidden_states = self.proj_act(hidden_states) hidden_states = self.proj_out(hidden_states) batch_size, channels, height, width = hidden_states.shape hidden_states = torch.cat( [ hidden_states, self.image_newline.reshape(1, channels, 1, 1) .expand(batch_size, channels, height, 1) .to(dtype, non_blocking=True), ], dim=-1, ) hidden_states = hidden_states.reshape(batch_size, channels, -1).permute(0, 2, 1) hidden_states = self.mlp(hidden_states) begin = ( self.image_begin.reshape(1, 1, -1) .expand(batch_size, 1, hidden_states.shape[-1]) .to(device=hidden_states.device, dtype=dtype, non_blocking=True) ) end = ( self.image_end.reshape(1, 1, -1) .expand(batch_size, 1, hidden_states.shape[-1]) .to(device=hidden_states.device, dtype=dtype, non_blocking=True) ) hidden_states = torch.cat([begin, hidden_states, end], dim=1) return self.after_rms(hidden_states) 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 HunYuanVLVisionAttention(nn.Module): def __init__(self, config: HunYuanVLVisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.num_heads = config.attention_heads self.head_dim = getattr(config, "head_dim", self.embed_dim // self.num_heads) self.scaling = self.head_dim**-0.5 self.num_key_value_groups = 1 self.q_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True) self.k_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=True) self.is_causal = False @deprecate_kwarg("hidden_state", new_name="hidden_states", version="v5.20") def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None]: query = self.q_proj(hidden_states) key = self.k_proj(hidden_states) value = self.v_proj(hidden_states) batch_size, seq_len, _ = query.shape query_states = query.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) value_states = value.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) if is_flash_attention_requested(self.config): # Flash Attention: Use cu_seqlens for variable length attention max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, scaling=self.scaling, cu_seq_lens_q=cu_seqlens, cu_seq_lens_k=cu_seqlens, max_length_q=max_seqlen, max_length_k=max_seqlen, is_causal=False, **kwargs, ) else: # Other implementations: Process each chunk separately lengths = cu_seqlens[1:] - cu_seqlens[:-1] splits = [ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) ] attn_outputs = [ attention_interface( self, q, k, v, scaling=self.scaling, is_causal=False, **kwargs, )[0] for q, k, v in zip(*splits) ] attn_output = torch.cat(attn_outputs, dim=1) attn_weights = None attn_output = attn_output.reshape(batch_size, seq_len, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class HunYuanVLVisionBlock(GradientCheckpointingLayer): def __init__(self, config: HunYuanVLVisionConfig): super().__init__() self.embed_dim = config.hidden_size self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.self_attn = HunYuanVLVisionAttention(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = HunYuanVLVisionMLP(config) @auto_docstring def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.FloatTensor: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states 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_func_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 apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): """ Apply HunYuan's multimodal rotary embedding to ``q`` and ``k``. `mrope_section` partitions half of the attention head dimension across the multimodal axes produced by `HunYuanVLModel.get_rope_index`. The section order matches the position-id channel order: `(width, height, image_index)` for 3-axis multimodal RoPE and `(position, width, height, image_index)` for 4-axis multimodal RoPE. """ x_dim = len(mrope_section) mrope_section = [int(section) * 2 for section in mrope_section] if sum(mrope_section) != cos.shape[-1]: raise ValueError( f"Illegal partition for multimodal RoPE: expected {cos.shape[-1]} rotary dims, got {sum(mrope_section)}" ) cos = torch.cat([m[i % x_dim] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1) sin = torch.cat([m[i % x_dim] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1) cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) origin_dtype = q.dtype q, k = q.float(), k.float() cos, sin = cos.float(), sin.float() q_out = (q * cos) + (rotate_half(q) * sin) k_out = (k * cos) + (rotate_half(k) * sin) return q_out.to(origin_dtype), k_out.to(origin_dtype) @use_kernelized_func(apply_rotary_pos_emb) class HunYuanVLDenseV1Attention(nn.Module): """ HunYuan dense attention with optional multimodal rotary embeddings. When ``rope_parameters['mrope_section']`` is set, queries and keys are rotated with HunYuan's multimodal axes. During decoding all axes carry the same text-only positions, reducing the operation to standard 1D RoPE. """ def __init__(self, config: HunYuanVLTextConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) self.query_layernorm = HunYuanVLRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.key_layernorm = HunYuanVLRMSNorm(self.head_dim, eps=config.rms_norm_eps) rope_parameters = getattr(config, "rope_parameters", None) or {} self.mrope_section = rope_parameters.get("mrope_section") def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( query_states, key_states, cos, sin, self.mrope_section ) query_states = self.query_layernorm(query_states) key_states = self.key_layernorm(key_states) 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(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class HunYuanVLMLP(nn.Module): def __init__(self, config: HunYuanVLConfig, layer_idx=None, is_shared_mlp=False): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] self.layer_idx = layer_idx def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class HunYuanVLDenseV1DecoderLayer(GradientCheckpointingLayer): def __init__(self, config: HunYuanVLTextConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = HunYuanVLDenseV1Attention(config=config, layer_idx=layer_idx) self.mlp = HunYuanVLMLP(config) self.input_layernorm = HunYuanVLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = HunYuanVLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.layer_idx = layer_idx def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) 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 = residual + hidden_states return hidden_states @auto_docstring class HunYuanVLPreTrainedModel(PreTrainedModel): config: HunYuanVLConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["HunYuanVLDenseV1DecoderLayer", "HunYuanVLVisionBlock"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": HunYuanVLDenseV1DecoderLayer, "attentions": HunYuanVLDenseV1Attention, } input_modalities = ("image", "text") @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) # DynamicNTKAlphaRotary - unique to this model if "RotaryEmbedding" in module.__class__.__name__ and hasattr(module, "original_inv_freq"): if module.rope_type == "dynamic" and module.config.rope_parameters.get("alpha"): dim = module.config.head_dim rope_theta = module.config.rope_parameters["rope_theta"] alpha = module.config.rope_parameters["alpha"] base = rope_theta * alpha ** (dim / (dim - 2)) buffer_value = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) else: rope_fn = ( ROPE_INIT_FUNCTIONS[module.rope_type] if module.rope_type != "default" else module.compute_default_rope_parameters ) buffer_value, _ = rope_fn(module.config) init.copy_(module.inv_freq, buffer_value) init.copy_(module.original_inv_freq, buffer_value) if isinstance(module, HunYuanVLVisionPatchMerger): embed_std = module.config.text_hidden_size**-0.5 init.normal_(module.image_newline, mean=0.0, std=embed_std) init.normal_(module.image_begin, mean=0.0, std=embed_std) init.normal_(module.image_end, mean=0.0, std=embed_std) init.normal_(module.image_sep, mean=0.0, std=embed_std) class HunYuanVLVisionTransformer(HunYuanVLPreTrainedModel): """ HunYuanVL vision tower: patch embedding -> transformer blocks -> per-image patch merger. Inputs are flat per-patch pixel tensors plus an ``image_grid_thw`` tensor describing the spatial layout of every image in the batch. The output is the concatenation of merged image embeddings, ready to be scattered into the language-model embedding stream. """ config: HunYuanVLVisionConfig main_input_name = "pixel_values" input_modalities = ("image",) _can_record_outputs = { "hidden_states": HunYuanVLVisionBlock, "attentions": HunYuanVLVisionAttention, } def __init__(self, config: HunYuanVLVisionConfig): super().__init__(config) self.embeddings = HunYuanVLVisionPatchEmbed(config) self.layers = nn.ModuleList([HunYuanVLVisionBlock(config) for _ in range(config.num_hidden_layers)]) self.patch_merger = HunYuanVLVisionPatchMerger(config) self.gradient_checkpointing = False self.post_init() @capture_outputs @auto_docstring def forward( self, pixel_values: torch.Tensor, grid_thw: torch.LongTensor, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" pixel_values (`torch.Tensor` of shape `(num_patches, num_channels * patch_size * patch_size)`): Flat per-patch pixel features produced by the image processor. grid_thw (`torch.LongTensor` of shape `(num_images, 3)`): The temporal, height and width dimensions for each image. Each row contains `[t, h, w]` patch counts. """ hidden_states = self.embeddings(pixel_values, grid_thw) cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs) for layer in self.layers: hidden_states = layer(hidden_states, cu_seqlens=cu_seqlens, attention_mask=None, **kwargs) split_sizes = grid_thw.prod(dim=-1).tolist() split_items = torch.split(hidden_states, split_sizes, dim=1) processed_items = [] for grid, item in zip(grid_thw.tolist(), split_items): _, h, w = grid processed_items.append(self.patch_merger(item, size=(h, w))) image_features = torch.cat(processed_items, dim=1) return BaseModelOutputWithPooling( last_hidden_state=hidden_states, pooler_output=image_features, ) @auto_docstring class HunYuanVLTextModel(HunYuanVLPreTrainedModel): """Dense text backbone used inside [`HunYuanVLModel`].""" config: HunYuanVLTextConfig input_modalities = ("text",) def __init__(self, config: HunYuanVLTextConfig): 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( [HunYuanVLDenseV1DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = HunYuanVLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = HunYuanVLRotaryEmbedding(config=config) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring 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[TransformersKwargs], ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) num_mrope_axes = len(self.rotary_emb.mrope_section or []) # Expand to 3D with `num_mrope_axes` as first dim, if not yet expanded 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(num_mrope_axes, inputs_embeds.shape[0], -1) elif position_ids.ndim == 2: position_ids = position_ids[None, ...].expand(num_mrope_axes, position_ids.shape[0], -1) if position_ids.ndim == 3 and position_ids.shape[0] == num_mrope_axes + 1: text_position_ids = position_ids[0] else: # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids text_position_ids = None if num_mrope_axes and position_ids.dim() == 3: if position_ids.shape[0] != num_mrope_axes: raise ValueError( f"Expected {num_mrope_axes} multimodal RoPE channels, got " f"position_ids with shape {tuple(position_ids.shape)}." ) causal_mask = create_causal_mask( self.config, inputs_embeds, attention_mask, past_key_values=past_key_values, position_ids=text_position_ids, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) for decoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=text_position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) @auto_docstring( custom_intro=""" The HunYuanVL model which consists of a vision backbone and a language model, without a language modeling head. """ ) class HunYuanVLModel(HunYuanVLPreTrainedModel): base_model_prefix = "model" # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False def __init__(self, config: HunYuanVLConfig): super().__init__(config) self.language_model = HunYuanVLTextModel._from_config(config.text_config) self.rope_deltas = None # cache rope_deltas here self.vision_tower = HunYuanVLVisionTransformer(config.vision_config) # Initialize weights and apply final processing self.post_init() def get_vision_position_ids( self, grid_hw: list[int, int] | torch.Tensor, spatial_merge_size: int = 1, device: str | torch.device | None = None, ): """ Compute HunYuanVL multimodal RoPE spatial indices for the pooled image-token grid of a single image. The vision merger appends one newline-style token per image row, so the width channel spans `patch_w + 1` positions while the height channel repeats each row id over that extra column. """ grid_h, grid_w = (int(value) for value in grid_hw) llm_grid_h = grid_h // spatial_merge_size llm_grid_w = grid_w // spatial_merge_size position_height, position_width = torch.meshgrid( torch.arange(llm_grid_h, dtype=torch.long, device=device), torch.arange(llm_grid_w + 1, dtype=torch.long, device=device), indexing="ij", ) return torch.stack([position_width.flatten(), position_height.flatten()], dim=0) def get_rope_index( self, input_ids: torch.LongTensor, mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, ) -> tuple[torch.LongTensor, torch.LongTensor]: """ Build HunYuanVL multimodal RoPE position ids. `rope_parameters["mrope_section"]` controls both the number and order of multimodal RoPE axes. Each section value is the number of half-rotary dimensions assigned to the corresponding axis, and the sections must sum to `head_dim // 2`. This method returns only the multimodal rotary axes consumed by the text backbone. The last three axes are `(width, height, image_index)`; any preceding axes keep their default 1D sequence positions. `width` and `height` index the pooled visual grid inside one image. `image_index` is the ordinal of the image/frame in the input sequence, so all visual tokens from the first image get `0`, the second image get `1`, and so on. Text-only 1D positions for the causal mask are inferred by the text backbone and are not part of this return value. """ rope_parameters = self.config.text_config.rope_parameters or {} num_mrope_axes = len(rope_parameters.get("mrope_section", [])) if num_mrope_axes < 3: raise ValueError(f"HunYuanVL expects at least 3 multimodal RoPE axes, got {num_mrope_axes}.") position_ids = torch.zeros( num_mrope_axes, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) grid_iter = iter(image_grid_thw) if image_grid_thw is not None else None rope_deltas = [] image_index = 0 for batch_idx, current_input_ids in enumerate(input_ids): input_token_type = mm_token_type_ids[batch_idx] valid_token_mask = None if attention_mask is not None: valid_token_mask = attention_mask[batch_idx].bool() current_input_ids = current_input_ids[valid_token_mask] input_token_type = input_token_type[valid_token_mask] current_position_ids = torch.arange( current_input_ids.shape[-1], dtype=input_ids.dtype, device=input_ids.device ) current_position_ids = current_position_ids.view(1, -1).expand(num_mrope_axes, -1).clone() if grid_iter is not None: for modality_type, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): if modality_type != 1: continue group = list(group) span_start = group[0][0] span_end = group[-1][0] + 1 try: grid_thw = next(grid_iter) except StopIteration as error: raise ValueError( "Found more image placeholder spans than entries in `image_grid_thw`." ) from error vision_position_ids = self.get_vision_position_ids( grid_thw[1:], spatial_merge_size=self.config.vision_config.spatial_merge_size, device=input_ids.device, ) grid_tokens = vision_position_ids.shape[1] span_length = span_end - span_start if span_length == grid_tokens + 2: grid_start = span_start + 1 elif span_length == grid_tokens: grid_start = span_start else: raise ValueError( "Image placeholder span length does not match `image_grid_thw`: " f"span_length={span_length}, expected {grid_tokens} or {grid_tokens + 2}." ) grid_end = grid_start + grid_tokens offset = num_mrope_axes - 3 current_position_ids[offset : offset + 2, grid_start:grid_end] = vision_position_ids.to( dtype=input_ids.dtype ) current_position_ids[offset + 2, grid_start:grid_end] = image_index image_index += 1 if valid_token_mask is not None: position_ids[:, batch_idx, valid_token_mask] = current_position_ids else: position_ids[:, batch_idx] = current_position_ids rope_deltas.append(current_position_ids.max() + 1 - len(current_input_ids)) if image_grid_thw is not None and image_index != len(image_grid_thw): raise ValueError( "Found fewer image placeholder spans than entries in `image_grid_thw`: " f"spans={image_index}, images={len(image_grid_thw)}." ) rope_deltas = torch.tensor(rope_deltas, device=input_ids.device).unsqueeze(1) return position_ids, rope_deltas @accepts_precomputed_kwargs(modality="image") @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" Encode images and return the complete vision tower outputs. pixel_values (`torch.FloatTensor`): Flat per-patch pixel features produced by the image processor. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ vision_dtype = next(self.vision_tower.parameters()).dtype pixel_values = pixel_values.to(vision_dtype) return self.vision_tower(pixel_values, grid_thw=image_grid_thw, **kwargs) def get_placeholder_mask( self, input_ids: torch.LongTensor | None, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor | None = None, ) -> torch.BoolTensor: """ Compute a boolean mask over ``inputs_embeds`` selecting the positions that hold the visual placeholder token, and validate that the placeholder count matches the number of provided image features. """ if input_ids is None: placeholder_token_embed = self.get_input_embeddings()( torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_image_mask = (inputs_embeds == placeholder_token_embed).all(-1) else: special_image_mask = input_ids == self.config.image_token_id n_image_tokens = special_image_mask.sum() special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel(): raise ValueError( f"Image features and image tokens do not match: tokens: {n_image_tokens}, " f"features {image_features.shape[0]}" ) return special_image_mask def compute_3d_position_ids( self, input_ids: torch.LongTensor | None, image_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.LongTensor | None: past_seen_tokens = 0 if past_key_values is None else past_key_values.get_seq_length() if image_grid_thw is not None and mm_token_type_ids is None and input_ids is not None: raise ValueError( "Multimodal data was passed via `image_grid_thw` but `mm_token_type_ids` is missing. Please pass " "`mm_token_type_ids` to the model so that multimodal RoPE can be computed correctly. " "`mm_token_type_ids` is returned by the processor alongside `input_ids`." ) if input_ids is None or image_grid_thw is None or past_seen_tokens != 0: return None rope_positions, rope_deltas = self.get_rope_index( input_ids, mm_token_type_ids=mm_token_type_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask, ) self.rope_deltas = rope_deltas return rope_positions @deprecate_kwarg("rope_deltas", version="v5.10") @can_return_tuple @auto_docstring 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, pixel_values: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> HunYuanVLModelOutputWithPast: r""" pixel_values (`torch.FloatTensor`, *optional*): Flat per-patch pixel features produced by the image processor. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) image_embeds = None if pixel_values is not None and image_grid_thw is not None: image_outputs = self.get_image_features(pixel_values, image_grid_thw, return_dict=True) image_embeds = image_outputs.pooler_output image_embeds = image_embeds.to(inputs_embeds.device, dtype=inputs_embeds.dtype, non_blocking=True) image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds, ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if position_ids is None: position_ids = self.compute_3d_position_ids( input_ids=input_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask, past_key_values=past_key_values, mm_token_type_ids=mm_token_type_ids, ) outputs: BaseModelOutputWithPast = self.language_model( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) return HunYuanVLModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_embeds if pixel_values is not None else None, ) @auto_docstring class HunYuanVLForConditionalGeneration(HunYuanVLPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: HunYuanVLConfig def __init__(self, config: HunYuanVLConfig): super().__init__(config) self.model = HunYuanVLModel(config) self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) self.vocab_size = config.text_config.vocab_size self.post_init() def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs) @can_return_tuple @auto_docstring 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, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, pixel_values: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: r""" pixel_values (`torch.FloatTensor`, *optional*): Flat per-patch pixel features produced by the image processor. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. Example: ```python >>> from transformers import AutoProcessor, HunYuanVLForConditionalGeneration >>> import torch >>> model_id = "tencent/HunyuanOCR" >>> processor = AutoProcessor.from_pretrained(model_id, backend="pil") >>> model = HunYuanVLForConditionalGeneration.from_pretrained( ... model_id, attn_implementation="eager", torch_dtype=torch.bfloat16, device_map="auto" ... ) >>> messages = [ ... { ... "role": "user", ... "content": [ ... {"type": "image", "image": "path/to/your/image.jpg"}, ... {"type": "text", "text": "Extract the text from the image."}, ... ], ... } ... ] >>> inputs = processor.apply_chat_template( ... messages, ... tokenize=True, ... add_generation_prompt=True, ... return_tensors="pt", ... return_dict=True, ... processor_kwargs={"padding": True}, ... ) >>> with torch.no_grad(): ... generated_ids = model.generate(**inputs, max_new_tokens=128) >>> generated_trimmed = generated_ids[0][inputs["input_ids"].shape[-1]:] >>> print(processor.decode(generated_trimmed, skip_special_tokens=True)) ```""" outputs: HunYuanVLModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, pixel_values=pixel_values, image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, **kwargs, ) hidden_states = outputs.last_hidden_state 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.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, past_key_values=None, attention_mask=None, inputs_embeds=None, position_ids=None, use_cache=True, pixel_values=None, image_grid_thw=None, mm_token_type_ids=None, is_first_iteration=False, **kwargs, ): model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, position_ids=position_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs, ) position_ids = model_inputs.get("position_ids") if position_ids is not None and position_ids.dim() == 3: rope_parameters = self.config.text_config.rope_parameters or {} num_mrope_axes = len(rope_parameters.get("mrope_section", [])) if num_mrope_axes and position_ids.shape[0] != num_mrope_axes: batch_size_source = model_inputs.get("input_ids") if batch_size_source is None: batch_size_source = model_inputs.get("inputs_embeds") if batch_size_source is None: batch_size_source = attention_mask batch_size = batch_size_source.shape[0] if position_ids.shape[0] % num_mrope_axes == 0 and position_ids.shape[1] != batch_size: expand_size = position_ids.shape[0] // num_mrope_axes position_ids = position_ids.view(expand_size, num_mrope_axes, *position_ids.shape[1:]) position_ids = position_ids.permute(1, 0, 2, 3).reshape(num_mrope_axes, -1, position_ids.shape[-1]) model_inputs["position_ids"] = position_ids if not is_first_iteration and use_cache: model_inputs["pixel_values"] = None return model_inputs def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): # Same as qwen-vl with variable `num_mrope_axes` based on config values text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs) rope_parameters = self.config.text_config.rope_parameters or {} num_mrope_axes = len(rope_parameters.get("mrope_section", [])) past_length = 0 if (cache := model_kwargs.get("past_key_values")) is not None: past_length = cache.get_seq_length() if past_length != 0 and self.model.rope_deltas is not None: return (text_positions + self.model.rope_deltas).unsqueeze(0).expand(num_mrope_axes, -1, -1) if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] if ( is_input_ids and model_kwargs.get("mm_token_type_ids") is not None and model_kwargs.get("image_grid_thw") is not None ): rope_positions, rope_deltas = self.model.get_rope_index( inputs_tensor, mm_token_type_ids=model_kwargs.get("mm_token_type_ids"), image_grid_thw=model_kwargs.get("image_grid_thw"), attention_mask=model_kwargs.get("attention_mask"), ) self.model.rope_deltas = rope_deltas else: rope_positions = text_positions.unsqueeze(0).expand(num_mrope_axes, -1, -1) self.model.rope_deltas = torch.zeros( inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device ) return rope_positions __all__ = [ "HunYuanVLPreTrainedModel", "HunYuanVLModel", "HunYuanVLTextModel", "HunYuanVLVisionTransformer", "HunYuanVLForConditionalGeneration", ]