Midea-AIRC/ECHO_block8
054
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_qwen2_5_vl.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.9#10# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX11# and OPT implementations in this library. It has been modified from its12# original forms to accommodate minor architectural differences compared13# to GPT-NeoX and OPT used by the Meta AI team that trained the model.14#15# Licensed under the Apache License, Version 2.0 (the "License");16# you may not use this file except in compliance with the License.17# You may obtain a copy of the License at18#19# http://www.apache.org/licenses/LICENSE-2.020#21# Unless required by applicable law or agreed to in writing, software22# distributed under the License is distributed on an "AS IS" BASIS,23# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.24# See the License for the specific language governing permissions and25# limitations under the License.26 27from dataclasses import dataclass28from typing import Any, Callable, Optional, Union29 30import torch31import torch.nn as nn32import torch.nn.functional as F33 34from transformers.activations import ACT2FN35from transformers.cache_utils import Cache, DynamicCache36from transformers.generation import GenerationMixin37from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask38from transformers.modeling_flash_attention_utils import FlashAttentionKwargs39from transformers.modeling_layers import GradientCheckpointingLayer40from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput41from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update42from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel43from transformers.processing_utils import Unpack44from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging45from .configuration_echo import EchoConfig, EchoTextConfig, EchoVisionConfig46 47# Add these imports48from typing import List, Tuple49from einops import rearrange50 51try:52 from flash_attn import flash_attn_func, flash_attn_varlen_func53 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input54except:55 pass56 57try:58 from liger_kernel.ops.swiglu import LigerSiLUMulFunction # noqa: F40159 liger_kernel_is_available = True60except ImportError:61 liger_kernel_is_available = False62 63from flash_attn.ops.triton.layer_norm import rms_norm_fn as flash_rms_norm64 65logger = logging.get_logger(__name__)66 67class EchoMLP(nn.Module):68 def __init__(self, config, bias: bool = False):69 super().__init__()70 self.hidden_size = config.hidden_size71 self.intermediate_size = config.intermediate_size72 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)73 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)74 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)75 self.act_fn = ACT2FN[config.hidden_act]76 77 def forward(self, hidden_state):78 return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))79 80 81class Echo_VisionPatchEmbed(nn.Module):82 def __init__(83 self,84 patch_size: int = 14,85 temporal_patch_size: int = 2,86 in_channels: int = 3,87 embed_dim: int = 1152,88 ) -> None:89 super().__init__()90 self.patch_size = patch_size91 self.temporal_patch_size = temporal_patch_size92 self.in_channels = in_channels93 self.embed_dim = embed_dim94 95 kernel_size = [temporal_patch_size, patch_size, patch_size]96 self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)97 98 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:99 target_dtype = self.proj.weight.dtype100 hidden_states = hidden_states.view(101 -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size102 )103 hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)104 return hidden_states105 106 107class Echo_VisionRotaryEmbedding(nn.Module):108 def __init__(self, dim: int, theta: float = 10000.0) -> None:109 super().__init__()110 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))111 self.register_buffer("inv_freq", inv_freq, persistent=False)112 113 def forward(self, seqlen: int) -> torch.Tensor:114 seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)115 freqs = torch.outer(seq, self.inv_freq)116 return freqs117 118 119class Qwen2RMSNorm(nn.Module):120 def __init__(self, hidden_size, eps=1e-6):121 """122 Qwen2RMSNorm is equivalent to T5LayerNorm123 """124 super().__init__()125 self.weight = nn.Parameter(torch.ones(hidden_size))126 self.variance_epsilon = eps127 128 def forward(self, hidden_states):129 input_dtype = hidden_states.dtype130 hidden_states = hidden_states.to(torch.float32)131 variance = hidden_states.pow(2).mean(-1, keepdim=True)132 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)133 return self.weight * hidden_states.to(input_dtype)134 135 def extra_repr(self):136 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"137 138 139class EchoPatchMerger(nn.Module):140 def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:141 super().__init__()142 self.hidden_size = context_dim * (spatial_merge_size**2)143 self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6)144 self.mlp = nn.Sequential(145 nn.Linear(self.hidden_size, self.hidden_size),146 nn.GELU(),147 nn.Linear(self.hidden_size, dim),148 )149 150 def forward(self, x: torch.Tensor) -> torch.Tensor:151 x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))152 return x153 154 155def rotate_half(x):156 """Rotates half the hidden dims of the input."""157 x1 = x[..., : x.shape[-1] // 2]158 x2 = x[..., x.shape[-1] // 2 :]159 return torch.cat((-x2, x1), dim=-1)160 161 162def apply_rotary_pos_emb_vision(163 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor164) -> tuple[torch.Tensor, torch.Tensor]:165 orig_q_dtype = q.dtype166 orig_k_dtype = k.dtype167 q, k = q.float(), k.float()168 cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()169 q_embed = (q * cos) + (rotate_half(q) * sin)170 k_embed = (k * cos) + (rotate_half(k) * sin)171 q_embed = q_embed.to(orig_q_dtype)172 k_embed = k_embed.to(orig_k_dtype)173 return q_embed, k_embed174 175 176def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:177 """178 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,179 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)180 """181 batch, num_key_value_heads, slen, head_dim = hidden_states.shape182 if n_rep == 1:183 return hidden_states184 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)185 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)186 187 188def eager_attention_forward(189 module: nn.Module,190 query: torch.Tensor,191 key: torch.Tensor,192 value: torch.Tensor,193 attention_mask: Optional[torch.Tensor],194 scaling: float,195 dropout: float = 0.0,196 **kwargs,197):198 key_states = repeat_kv(key, module.num_key_value_groups)199 value_states = repeat_kv(value, module.num_key_value_groups)200 201 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling202 if attention_mask is not None:203 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]204 attn_weights = attn_weights + causal_mask205 206 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)207 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)208 attn_output = torch.matmul(attn_weights, value_states)209 attn_output = attn_output.transpose(1, 2).contiguous()210 211 return attn_output, attn_weights212 213 214class EchoVisionAttention(nn.Module):215 def __init__(self, config: EchoVisionConfig) -> None:216 super().__init__()217 self.dim = config.hidden_size218 self.num_heads = config.num_heads219 self.head_dim = self.dim // self.num_heads220 self.num_key_value_groups = 1 # needed for eager attention221 self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)222 self.proj = nn.Linear(self.dim, self.dim)223 self.scaling = self.head_dim**-0.5224 self.config = config225 self.attention_dropout = 0.0226 self.is_causal = False227 228 def forward(229 self,230 hidden_states: torch.Tensor,231 cu_seqlens: torch.Tensor,232 rotary_pos_emb: Optional[torch.Tensor] = None,233 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,234 **kwargs,235 ) -> torch.Tensor:236 seq_length = hidden_states.shape[0]237 query_states, key_states, value_states = (238 self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)239 )240 if position_embeddings is None:241 logger.warning_once(242 "The attention layers in this model are transitioning from computing the RoPE embeddings internally "243 "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed "244 "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be "245 "removed and `position_embeddings` will be mandatory."246 )247 emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)248 cos = emb.cos()249 sin = emb.sin()250 else:251 cos, sin = position_embeddings252 query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)253 254 query_states = query_states.transpose(0, 1).unsqueeze(0)255 key_states = key_states.transpose(0, 1).unsqueeze(0)256 value_states = value_states.transpose(0, 1).unsqueeze(0)257 258 attention_interface: Callable = eager_attention_forward259 if self.config._attn_implementation != "eager":260 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]261 262 if self.config._attn_implementation == "flash_attention_2":263 # Flash Attention 2: Use cu_seqlens for variable length attention264 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()265 attn_output, _ = attention_interface(266 self,267 query_states,268 key_states,269 value_states,270 attention_mask=None,271 scaling=self.scaling,272 dropout=0.0 if not self.training else self.attention_dropout,273 cu_seq_lens_q=cu_seqlens,274 cu_seq_lens_k=cu_seqlens,275 max_length_q=max_seqlen,276 max_length_k=max_seqlen,277 is_causal=False,278 **kwargs,279 )280 else:281 # Other implementations: Process each chunk separately282 lengths = cu_seqlens[1:] - cu_seqlens[:-1]283 splits = [284 torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)285 ]286 287 attn_outputs = [288 attention_interface(289 self,290 q,291 k,292 v,293 attention_mask=None,294 scaling=self.scaling,295 dropout=0.0 if not self.training else self.attention_dropout,296 is_causal=False,297 **kwargs,298 )[0]299 for q, k, v in zip(*splits)300 ]301 attn_output = torch.cat(attn_outputs, dim=1)302 303 attn_output = attn_output.reshape(seq_length, -1).contiguous()304 attn_output = self.proj(attn_output)305 return attn_output306 307 308class EchoVisionBlock(GradientCheckpointingLayer):309 def __init__(self, config, attn_implementation: str = "sdpa") -> None:310 super().__init__()311 self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)312 self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)313 self.attn = EchoVisionAttention(config=config)314 self.mlp = EchoMLP(config, bias=True)315 316 def forward(317 self,318 hidden_states: torch.Tensor,319 cu_seqlens: torch.Tensor,320 rotary_pos_emb: Optional[torch.Tensor] = None,321 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,322 **kwargs,323 ) -> torch.Tensor:324 hidden_states = hidden_states + self.attn(325 self.norm1(hidden_states),326 cu_seqlens=cu_seqlens,327 rotary_pos_emb=rotary_pos_emb,328 position_embeddings=position_embeddings,329 **kwargs,330 )331 hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))332 return hidden_states333 334 335@auto_docstring336class EchoPreTrainedModel(PreTrainedModel):337 config: EchoConfig338 base_model_prefix = "model"339 supports_gradient_checkpointing = True340 _no_split_modules = ["EchoDecoderLayer", "EchoVisionBlock"]341 _skip_keys_device_placement = "past_key_values"342 _supports_flash_attn = True343 _supports_sdpa = True344 345 _can_compile_fullgraph = True346 _supports_attention_backend = True347 348 349class Echo_VisionTransformerPretrainedModel(EchoPreTrainedModel):350 config: EchoVisionConfig351 _no_split_modules = ["EchoVisionBlock"]352 353 def __init__(self, config, *inputs, **kwargs) -> None:354 super().__init__(config, *inputs, **kwargs)355 self.spatial_merge_size = config.spatial_merge_size356 self.patch_size = config.patch_size357 self.fullatt_block_indexes = config.fullatt_block_indexes358 self.window_size = config.window_size359 self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size360 361 self.patch_embed = Echo_VisionPatchEmbed(362 patch_size=config.patch_size,363 temporal_patch_size=config.temporal_patch_size,364 in_channels=config.in_channels,365 embed_dim=config.hidden_size,366 )367 368 head_dim = config.hidden_size // config.num_heads369 self.rotary_pos_emb = Echo_VisionRotaryEmbedding(head_dim // 2)370 371 self.blocks = nn.ModuleList([EchoVisionBlock(config) for _ in range(config.depth)])372 self.merger = EchoPatchMerger(373 dim=config.out_hidden_size,374 context_dim=config.hidden_size,375 spatial_merge_size=config.spatial_merge_size,376 )377 self.gradient_checkpointing = False378 379 def rot_pos_emb(self, grid_thw):380 pos_ids = []381 for t, h, w in grid_thw:382 hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)383 hpos_ids = hpos_ids.reshape(384 h // self.spatial_merge_size,385 self.spatial_merge_size,386 w // self.spatial_merge_size,387 self.spatial_merge_size,388 )389 hpos_ids = hpos_ids.permute(0, 2, 1, 3)390 hpos_ids = hpos_ids.flatten()391 392 wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)393 wpos_ids = wpos_ids.reshape(394 h // self.spatial_merge_size,395 self.spatial_merge_size,396 w // self.spatial_merge_size,397 self.spatial_merge_size,398 )399 wpos_ids = wpos_ids.permute(0, 2, 1, 3)400 wpos_ids = wpos_ids.flatten()401 pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))402 pos_ids = torch.cat(pos_ids, dim=0)403 max_grid_size = grid_thw[:, 1:].max()404 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)405 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)406 return rotary_pos_emb407 408 def get_window_index(self, grid_thw):409 window_index: list = []410 cu_window_seqlens: list = [0]411 window_index_id = 0412 vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size413 414 for grid_t, grid_h, grid_w in grid_thw:415 llm_grid_h, llm_grid_w = (416 grid_h // self.spatial_merge_size,417 grid_w // self.spatial_merge_size,418 )419 index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)420 pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size421 pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size422 num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size423 num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size424 index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)425 index_padded = index_padded.reshape(426 grid_t,427 num_windows_h,428 vit_merger_window_size,429 num_windows_w,430 vit_merger_window_size,431 )432 index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(433 grid_t,434 num_windows_h * num_windows_w,435 vit_merger_window_size,436 vit_merger_window_size,437 )438 seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)439 index_padded = index_padded.reshape(-1)440 index_new = index_padded[index_padded != -100]441 window_index.append(index_new + window_index_id)442 cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]443 cu_window_seqlens.extend(cu_seqlens_tmp.tolist())444 window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()445 window_index = torch.cat(window_index, dim=0)446 447 return window_index, cu_window_seqlens448 449 def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:450 """451 Args:452 hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):453 The final hidden states of the model.454 grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):455 The temporal, height and width of feature shape of each image in LLM.456 457 Returns:458 `torch.Tensor`: hidden_states.459 """460 hidden_states = self.patch_embed(hidden_states)461 rotary_pos_emb = self.rot_pos_emb(grid_thw)462 window_index, cu_window_seqlens = self.get_window_index(grid_thw)463 cu_window_seqlens = torch.tensor(464 cu_window_seqlens,465 device=hidden_states.device,466 dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,467 )468 cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)469 470 seq_len, _ = hidden_states.size()471 hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)472 hidden_states = hidden_states[window_index, :, :]473 hidden_states = hidden_states.reshape(seq_len, -1)474 rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)475 rotary_pos_emb = rotary_pos_emb[window_index, :, :]476 rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)477 emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)478 position_embeddings = (emb.cos(), emb.sin())479 480 cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(481 dim=0,482 # Select dtype based on the following factors:483 # - FA2 requires that cu_seqlens_q must have dtype int32484 # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw485 # See https://github.com/huggingface/transformers/pull/34852 for more information486 dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,487 )488 cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)489 490 for layer_num, blk in enumerate(self.blocks):491 if layer_num in self.fullatt_block_indexes:492 cu_seqlens_now = cu_seqlens493 else:494 cu_seqlens_now = cu_window_seqlens495 496 hidden_states = blk(497 hidden_states,498 cu_seqlens=cu_seqlens_now,499 position_embeddings=position_embeddings,500 **kwargs,501 )502 503 hidden_states = self.merger(hidden_states)504 reverse_indices = torch.argsort(window_index)505 hidden_states = hidden_states[reverse_indices, :]506 507 return hidden_states508 509 510@dataclass511@auto_docstring(512 custom_intro="""513 Base class for Llava outputs, with hidden states and attentions.514 """515)516class EchoModelOutputWithPast(ModelOutput):517 r"""518 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):519 Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape520 `(batch_size, num_heads, sequence_length, embed_size_per_head)`)521 522 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see523 `past_key_values` input) to speed up sequential decoding.524 rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):525 The rope index difference between sequence length and multimodal rope.526 """527 528 last_hidden_state: torch.FloatTensor = None529 past_key_values: Optional[list[torch.FloatTensor]] = None530 hidden_states: Optional[tuple[torch.FloatTensor]] = None531 attentions: Optional[tuple[torch.FloatTensor]] = None532 rope_deltas: Optional[torch.LongTensor] = None533 534 535class EchoRotaryEmbedding(nn.Module):536 def __init__(self, config: EchoTextConfig, device=None):537 super().__init__()538 # BC: "rope_type" was originally "type"539 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:540 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))541 else:542 self.rope_type = "default"543 self.max_seq_len_cached = config.max_position_embeddings544 self.original_max_seq_len = config.max_position_embeddings545 546 self.config = config547 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]548 549 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)550 self.register_buffer("inv_freq", inv_freq, persistent=False)551 self.original_inv_freq = self.inv_freq552 553 @torch.no_grad()554 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)555 def forward(self, x, position_ids):556 # In contrast to other models, Echo has different position ids for the grids557 # So we expand the inv_freq to shape (3, ...)558 inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)559 position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)560 561 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"562 with torch.autocast(device_type=device_type, enabled=False): # Force float32563 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)564 emb = torch.cat((freqs, freqs), dim=-1)565 cos = emb.cos() * self.attention_scaling566 sin = emb.sin() * self.attention_scaling567 568 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)569 570 571class Qwen2MLP(nn.Module):572 def __init__(self, config):573 super().__init__()574 self.config = config575 self.hidden_size = config.hidden_size576 self.intermediate_size = config.intermediate_size577 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)578 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)579 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)580 self.act_fn = ACT2FN[config.hidden_act]581 582 def forward(self, x):583 if liger_kernel_is_available:584 return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))585 else:586 down_proj = self.down_proj(self.act_fn(587 self.gate_proj(x)) * self.up_proj(x))588 return down_proj589 590 591def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):592 """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).593 594 Explanation:595 Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding596 sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For597 vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.598 Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.599 For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,600 height and width) of text embedding is always the same, so the text embedding rotary position embedding has no601 difference with modern LLMs.602 603 Args:604 q (`torch.Tensor`): The query tensor.605 k (`torch.Tensor`): The key tensor.606 cos (`torch.Tensor`): The cosine part of the rotary embedding.607 sin (`torch.Tensor`): The sine part of the rotary embedding.608 position_ids (`torch.Tensor`):609 The position indices of the tokens corresponding to the query and key tensors. For example, this can be610 used to pass offsetted position ids when working with a KV-cache.611 mrope_section(`List(int)`):612 Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.613 unsqueeze_dim (`int`, *optional*, defaults to 1):614 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and615 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note616 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and617 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes618 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have619 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.620 Returns:621 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.622 """623 mrope_section = mrope_section * 2624 cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(625 unsqueeze_dim626 )627 sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(628 unsqueeze_dim629 )630 631 q_embed = (q * cos) + (rotate_half(q) * sin)632 k_embed = (k * cos) + (rotate_half(k) * sin)633 return q_embed, k_embed634 635 636class EchoAttention(nn.Module):637 """638 Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer639 and "Generating Long Sequences with Sparse Transformers".640 """641 642 def __init__(self, config: EchoTextConfig, layer_idx: Optional[int] = None):643 super().__init__()644 self.config = config645 self.layer_idx = layer_idx646 if layer_idx is None:647 logger.warning_once(648 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "649 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "650 "when creating this class."651 )652 653 self.hidden_size = config.hidden_size654 self.num_heads = config.num_attention_heads655 self.head_dim = self.hidden_size // self.num_heads656 self.num_key_value_heads = config.num_key_value_heads657 self.num_key_value_groups = self.num_heads // self.num_key_value_heads658 self.is_causal = True659 self.attention_dropout = config.attention_dropout660 self.rope_scaling = config.rope_scaling661 self.scaling = self.head_dim**-0.5662 663 if (self.head_dim * self.num_heads) != self.hidden_size:664 raise ValueError(665 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"666 f" and `num_heads`: {self.num_heads})."667 )668 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)669 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)670 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)671 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)672 self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None673 674 self.rotary_emb = EchoRotaryEmbedding(config=config)675 676 def forward(677 self,678 hidden_states: torch.Tensor,679 attention_mask: Optional[torch.Tensor] = None,680 position_ids: Optional[torch.LongTensor] = None,681 past_key_values: Optional[Cache] = None,682 output_attentions: bool = False,683 use_cache: bool = False,684 cache_position: Optional[torch.LongTensor] = None,685 store_kv: Optional[bool] = False,686 store_kv_len: Optional[int] = None,687 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC688 **kwargs: Unpack[FlashAttentionKwargs],689 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:690 input_shape = hidden_states.shape[:-1]691 hidden_shape = (*input_shape, -1, self.head_dim)692 bsz, q_len, _ = hidden_states.size()693 694 query_states = self.q_proj(hidden_states)695 key_states = self.k_proj(hidden_states)696 value_states = self.v_proj(hidden_states)697 698 query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)699 key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)700 value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)701 702 cos, sin = position_embeddings703 query_states, key_states = apply_multimodal_rotary_pos_emb(704 query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]705 )706 707 708 if past_key_values is not None and store_kv:709 if store_kv_len is not None and 0 < store_kv_len < q_len:710 # Partial-store: only cache KV for the first store_kv_len tokens (prev finalized block).711 # The remaining tokens (current masked block) are only used for this forward's attention.712 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position[:store_kv_len]}713 cached_k, cached_v = past_key_values.update(714 key_states[:, :, :store_kv_len, :],715 value_states[:, :, :store_kv_len, :],716 self.layer_idx,717 cache_kwargs,718 )719 # Attention sees: past KV (now including stored portion) + non-stored tokens720 key_states = torch.cat([cached_k, key_states[:, :, store_kv_len:, :]], dim=-2)721 value_states = torch.cat([cached_v, value_states[:, :, store_kv_len:, :]], dim=-2)722 else:723 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}724 key_states, value_states = past_key_values.update(725 key_states, value_states, self.layer_idx, cache_kwargs726 )727 728 elif past_key_values is not None and (not store_kv) and len(past_key_values) > self.layer_idx:729 past_key_states, past_value_states = past_key_values[self.layer_idx]730 key_states = torch.cat([past_key_states, key_states], dim=-2)731 value_states = torch.cat([past_value_states, value_states], dim=-2)732 733 734 attention_interface: Callable = eager_attention_forward735 if self.config._attn_implementation != "eager":736 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]737 738 attention_mask = attention_mask.bool() if attention_mask is not None else None739 attn_weights = None740 if torch.all(attention_mask): # decoding741 query_states = query_states.transpose(1, 2)742 key_states = key_states.transpose(1, 2)743 value_states = value_states.transpose(1, 2)744 attn_output = flash_attn_func(745 query_states,746 key_states,747 value_states,748 causal=False,749 softmax_scale=self.scaling750 )751 attn_output = rearrange(attn_output, 'b l h d -> b l (h d)')752 else: # prefilling753 attn_output = F.scaled_dot_product_attention(754 query=query_states,755 key=key_states,756 value=value_states,757 attn_mask=attention_mask,758 is_causal=False,759 scale=self.scaling,760 enable_gqa=True761 )762 attn_output = rearrange(attn_output, 'b h l d -> b l (h d)')763 764 attn_output = attn_output.reshape(*input_shape, -1).contiguous()765 attn_output = self.o_proj(attn_output)766 return attn_output, attn_weights767 768 769class EchoDecoderLayer(GradientCheckpointingLayer):770 def __init__(self, config: EchoTextConfig, layer_idx: int):771 super().__init__()772 self.hidden_size = config.hidden_size773 774 if config.use_sliding_window and config._attn_implementation != "flash_attention_2":775 logger.warning_once(776 f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "777 "unexpected results may be encountered."778 )779 self.self_attn = EchoAttention(config, layer_idx)780 781 self.mlp = Qwen2MLP(config)782 self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)783 self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)784 self.attention_type = config.layer_types[layer_idx]785 786 def forward(787 self,788 hidden_states: torch.Tensor,789 attention_mask: Optional[torch.Tensor] = None,790 position_ids: Optional[torch.LongTensor] = None,791 past_key_values: Optional[tuple[torch.Tensor]] = None,792 output_attentions: Optional[bool] = False,793 use_cache: Optional[bool] = False,794 cache_position: Optional[torch.LongTensor] = None,795 store_kv: Optional[bool] = False,796 store_kv_len: Optional[int] = None,797 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC798 **kwargs: Unpack[FlashAttentionKwargs],799 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:800 """801 Args:802 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`803 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size804 `(batch, sequence_length)` where padding elements are indicated by 0.805 output_attentions (`bool`, *optional*):806 Whether or not to return the attentions tensors of all attention layers. See `attentions` under807 returned tensors for more detail.808 use_cache (`bool`, *optional*):809 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding810 (see `past_key_values`).811 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states812 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):813 Indices depicting the position of the input sequence tokens in the sequence.814 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):815 Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,816 with `head_dim` being the embedding dimension of each attention head.817 kwargs (`dict`, *optional*):818 Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code819 into the model820 """821 822 residual = hidden_states823 824 hidden_states = self.input_layernorm(hidden_states)825 826 # Self Attention827 hidden_states, self_attn_weights = self.self_attn(828 hidden_states=hidden_states,829 attention_mask=attention_mask,830 position_ids=position_ids,831 past_key_values=past_key_values,832 output_attentions=output_attentions,833 use_cache=use_cache,834 cache_position=cache_position,835 store_kv=store_kv,836 store_kv_len=store_kv_len,837 position_embeddings=position_embeddings,838 **kwargs,839 )840 hidden_states = residual + hidden_states841 842 # Fully Connected843 residual = hidden_states844 hidden_states = self.post_attention_layernorm(hidden_states)845 hidden_states = self.mlp(hidden_states)846 hidden_states = residual + hidden_states847 848 outputs = (hidden_states,)849 850 if output_attentions:851 outputs += (self_attn_weights,)852 853 return outputs854 855 856@auto_docstring857class EchoTextModel(EchoPreTrainedModel):858 config: EchoTextConfig859 860 def __init__(self, config: EchoTextConfig):861 super().__init__(config)862 self.padding_idx = config.pad_token_id863 self.vocab_size = config.vocab_size864 865 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)866 self.layers = nn.ModuleList(867 [EchoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]868 )869 self._attn_implementation = config._attn_implementation870 self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)871 self.rotary_emb = EchoRotaryEmbedding(config=config)872 self.has_sliding_layers = "sliding_attention" in self.config.layer_types873 874 self.gradient_checkpointing = False875 # Initialize weights and apply final processing876 self.post_init()877 878 @auto_docstring879 def forward(880 self,881 input_ids: Optional[torch.LongTensor] = None,882 attention_mask: Optional[torch.Tensor] = None,883 position_ids: Optional[torch.LongTensor] = None,884 past_key_values: Optional[Cache] = None,885 inputs_embeds: Optional[torch.FloatTensor] = None,886 use_cache: Optional[bool] = None,887 output_attentions: Optional[bool] = None,888 output_hidden_states: Optional[bool] = None,889 return_dict: Optional[bool] = None,890 cache_position: Optional[torch.LongTensor] = None,891 store_kv: Optional[bool] = False,892 store_kv_len: Optional[int] = None,893 **kwargs: Unpack[FlashAttentionKwargs],894 ) -> Union[tuple, BaseModelOutputWithPast]:895 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions896 output_hidden_states = (897 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states898 )899 use_cache = use_cache if use_cache is not None else self.config.use_cache900 901 return_dict = return_dict if return_dict is not None else self.config.use_return_dict902 903 if (input_ids is None) ^ (inputs_embeds is not None):904 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")905 906 if self.gradient_checkpointing and self.training:907 if use_cache:908 logger.warning_once(909 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."910 )911 use_cache = False912 913 # torch.jit.trace() doesn't support cache objects in the output914 if use_cache and past_key_values is None and not torch.jit.is_tracing():915 past_key_values = DynamicCache()916 917 if inputs_embeds is None:918 inputs_embeds = self.embed_tokens(input_ids)919 920 if cache_position is None:921 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0922 cache_position = torch.arange(923 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device924 )925 926 # the hard coded `3` is for temporal, height and width.927 if position_ids is None:928 position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)929 elif position_ids.ndim == 2:930 position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)931 932 # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions933 # where each dim indicates visual spatial positions for temporal/height/width grids.934 # There are two scenarios when FA2-like packed masking might be activated.935 # 1. User specifically passed packed `position_ids` and no attention mask.936 # In this case we expect the useer to create correct position ids for all 3 grids937 # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]938 # 2. User runs forward with no attention mask and no position ids. In this case, position ids939 # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are940 # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass941 # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`942 if position_ids.ndim == 3 and position_ids.shape[0] == 4:943 text_position_ids = position_ids[0]944 position_ids = position_ids[1:]945 else:946 text_position_ids = position_ids[0]947 948 # It may already have been prepared by e.g. `generate`949 if not isinstance(causal_mask_mapping := attention_mask, dict):950 # Prepare mask arguments951 mask_kwargs = {952 "config": self.config,953 "input_embeds": inputs_embeds,954 "attention_mask": attention_mask,955 "cache_position": cache_position,956 "past_key_values": past_key_values,957 "position_ids": text_position_ids,958 }959 # Create the masks960 causal_mask_mapping = {961 "full_attention": create_causal_mask(**mask_kwargs),962 }963 # The sliding window alternating layers are not always activated depending on the config964 if self.has_sliding_layers:965 causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)966 967 hidden_states = inputs_embeds968 969 # create position embeddings to be shared across the decoder layers970 position_embeddings = self.rotary_emb(hidden_states, position_ids)971 972 # decoder layers973 all_hidden_states = () if output_hidden_states else None974 all_self_attns = () if output_attentions else None975 976 for decoder_layer in self.layers:977 if output_hidden_states:978 all_hidden_states += (hidden_states,)979 980 layer_outputs = decoder_layer(981 hidden_states,982 attention_mask=causal_mask_mapping[decoder_layer.attention_type],983 position_ids=text_position_ids,984 past_key_values=past_key_values,985 output_attentions=output_attentions,986 use_cache=use_cache,987 cache_position=cache_position,988 position_embeddings=position_embeddings,989 store_kv=store_kv,990 store_kv_len=store_kv_len,991 **kwargs,992 )993 994 hidden_states = layer_outputs[0]995 996 if output_attentions:997 all_self_attns += (layer_outputs[1],)998 999 hidden_states = self.norm(hidden_states)1000 1001 # add hidden states from the last decoder layer1002 if output_hidden_states:1003 all_hidden_states += (hidden_states,)1004 1005 if not return_dict:1006 return tuple(1007 v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None1008 )1009 return BaseModelOutputWithPast(1010 last_hidden_state=hidden_states,1011 past_key_values=past_key_values,1012 hidden_states=all_hidden_states,1013 attentions=all_self_attns,1014 )1015 1016 1017@auto_docstring1018class EchoModel(EchoPreTrainedModel):1019 base_model_prefix = ""1020 _checkpoint_conversion_mapping = {"^model": "language_model"}1021 config: EchoConfig1022 _no_split_modules = ["EchoDecoderLayer", "EchoVisionBlock"]1023 1024 def __init__(self, config):1025 super().__init__(config)1026 self.visual = Echo_VisionTransformerPretrainedModel._from_config(config.vision_config)1027 self.language_model = EchoTextModel._from_config(config.text_config)1028 self.rope_deltas = None # cache rope_deltas here1029 1030 # Initialize weights and apply final processing1031 self.post_init()1032 1033 def get_input_embeddings(self):1034 return self.language_model.get_input_embeddings()1035 1036 def set_input_embeddings(self, value):1037 self.language_model.set_input_embeddings(value)1038 1039 def set_decoder(self, decoder):1040 self.language_model = decoder1041 1042 def get_decoder(self):1043 return self.language_model1044 1045 def get_rope_index(1046 self,1047 input_ids: Optional[torch.LongTensor] = None,1048 image_grid_thw: Optional[torch.LongTensor] = None,1049 video_grid_thw: Optional[torch.LongTensor] = None,1050 second_per_grid_ts: Optional[torch.Tensor] = None,1051 attention_mask: Optional[torch.Tensor] = None,1052 ) -> tuple[torch.Tensor, torch.Tensor]:1053 """1054 Calculate the 3D rope index based on image and video's temporal, height and width in LLM.1055 1056 Explanation:1057 Each embedding sequence contains vision embedding and text embedding or just contains text embedding.1058 1059 For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.1060 Examples:1061 input_ids: [T T T T T], here T is for text.1062 temporal position_ids: [0, 1, 2, 3, 4]1063 height position_ids: [0, 1, 2, 3, 4]1064 width position_ids: [0, 1, 2, 3, 4]1065 1066 For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part1067 and 1D rotary position embedding for text part.1068 Examples:1069 Temporal (Time): 3 patches, representing different segments of the video in time.1070 Height: 2 patches, dividing each frame vertically.1071 Width: 2 patches, dividing each frame horizontally.1072 We also have some important parameters:1073 fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.1074 tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.1075 temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.1076 interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.1077 input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.1078 vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]1079 vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]1080 vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]1081 text temporal position_ids: [101, 102, 103, 104, 105]1082 text height position_ids: [101, 102, 103, 104, 105]1083 text width position_ids: [101, 102, 103, 104, 105]1084 Here we calculate the text start position_ids as the max vision position_ids plus 1.1085 1086 Args:1087 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):1088 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide1089 it.1090 image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):1091 The temporal, height and width of feature shape of each image in LLM.1092 video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):1093 The temporal, height and width of feature shape of each video in LLM.1094 second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):1095 The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.1096 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1097 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:1098 1099 - 1 for tokens that are **not masked**,1100 - 0 for tokens that are **masked**.1101 1102 Returns:1103 position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)1104 mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)1105 """1106 spatial_merge_size = self.config.vision_config.spatial_merge_size1107 image_token_id = self.config.image_token_id1108 video_token_id = self.config.video_token_id1109 vision_start_token_id = self.config.vision_start_token_id1110 mrope_position_deltas = []1111 if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):1112 total_input_ids = input_ids1113 if attention_mask is None:1114 attention_mask = torch.ones_like(total_input_ids)1115 position_ids = torch.ones(1116 3,1117 input_ids.shape[0],1118 input_ids.shape[1],1119 dtype=input_ids.dtype,1120 device=input_ids.device,1121 )1122 image_index, video_index = 0, 01123 attention_mask = attention_mask.to(total_input_ids.device)1124 for i, input_ids in enumerate(total_input_ids):1125 input_ids = input_ids[attention_mask[i] == 1]1126 image_nums, video_nums = 0, 01127 vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)1128 vision_tokens = input_ids[vision_start_indices + 1]1129 image_nums = (vision_tokens == image_token_id).sum()1130 video_nums = (vision_tokens == video_token_id).sum()1131 input_tokens = input_ids.tolist()1132 llm_pos_ids_list: list = []1133 st = 01134 remain_images, remain_videos = image_nums, video_nums1135 for _ in range(image_nums + video_nums):1136 if image_token_id in input_tokens and remain_images > 0:1137 ed_image = input_tokens.index(image_token_id, st)1138 else:1139 ed_image = len(input_tokens) + 11140 if video_token_id in input_tokens and remain_videos > 0:1141 ed_video = input_tokens.index(video_token_id, st)1142 else:1143 ed_video = len(input_tokens) + 11144 if ed_image < ed_video:1145 t, h, w = (1146 image_grid_thw[image_index][0],1147 image_grid_thw[image_index][1],1148 image_grid_thw[image_index][2],1149 )1150 second_per_grid_t = 01151 image_index += 11152 remain_images -= 11153 ed = ed_image1154 1155 else:1156 t, h, w = (1157 video_grid_thw[video_index][0],1158 video_grid_thw[video_index][1],1159 video_grid_thw[video_index][2],1160 )1161 if second_per_grid_ts is not None:1162 second_per_grid_t = second_per_grid_ts[video_index]1163 else:1164 second_per_grid_t = 1.01165 video_index += 11166 remain_videos -= 11167 ed = ed_video1168 llm_grid_t, llm_grid_h, llm_grid_w = (1169 t.item(),1170 h.item() // spatial_merge_size,1171 w.item() // spatial_merge_size,1172 )1173 text_len = ed - st1174 1175 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 01176 llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)1177 1178 range_tensor = torch.arange(llm_grid_t).view(-1, 1)1179 expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)1180 1181 ## normalize type, send to device.1182 second_per_grid_t = torch.as_tensor(1183 second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device1184 )1185 1186 time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second1187 1188 time_tensor_long = time_tensor.long()1189 t_index = time_tensor_long.flatten()1190 1191 h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()1192 w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()1193 llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)1194 st = ed + llm_grid_t * llm_grid_h * llm_grid_w1195 1196 if st < len(input_tokens):1197 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 01198 text_len = len(input_tokens) - st1199 llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)1200