linyq/kiwi-edit-5b-reference-only-diffusers
441
1# coding=utf-82# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import base6416import copy17import math18import os19import sys20import time21import warnings22from concurrent.futures import ThreadPoolExecutor23from dataclasses import dataclass24from functools import lru_cache25from io import BytesIO26from typing import Any, Callable, Dict, List, Optional, Tuple, Union27 28import numpy as np29import requests30import torch31import torch.nn as nn32import torch.nn.functional as F33import torchvision34from packaging import version35from PIL import Image36from torchvision import io, transforms37from torchvision.transforms import InterpolationMode38 39from diffusers import ModelMixin, ConfigMixin40from diffusers.configuration_utils import register_to_config41from transformers.activations import ACT2FN42from transformers.cache_utils import Cache, DynamicCache43from transformers.generation import GenerationMixin44from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask45from transformers.modeling_flash_attention_utils import FlashAttentionKwargs46from transformers.modeling_layers import GradientCheckpointingLayer47from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput48from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update49from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel50from transformers.processing_utils import Unpack51from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging52from transformers.utils.deprecation import deprecate_kwarg53from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm54from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig55 56logger = logging.get_logger(__name__)57 58# ─────────────────────────────────────────────────────────────────────────────59# Vision processing utilities (formerly qwen_vl_utils.py)60# ─────────────────────────────────────────────────────────────────────────────61 62MAX_RATIO = 20063SPATIAL_MERGE_SIZE = 264IMAGE_MIN_TOKEN_NUM = 465IMAGE_MAX_TOKEN_NUM = 1638466VIDEO_MIN_TOKEN_NUM = 12867VIDEO_MAX_TOKEN_NUM = 76868 69FPS = 2.070FRAME_FACTOR = 271FPS_MIN_FRAMES = 472FPS_MAX_FRAMES = 1673MAX_NUM_WORKERS_FETCH_VIDEO = 874 75MODEL_SEQ_LEN = int(float(os.environ.get('MODEL_SEQ_LEN', 128000)))76 77 78# ─────────────────────────────────────────────────────────────────────────────79# Qwen2.5-VL model (formerly modeling_qwen2_5_vl.py)80# ─────────────────────────────────────────────────────────────────────────────81 82class Qwen2_5_VLMLP(nn.Module):83 def __init__(self, config, bias: bool = False):84 super().__init__()85 self.hidden_size = config.hidden_size86 self.intermediate_size = config.intermediate_size87 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)88 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)89 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)90 self.act_fn = ACT2FN[config.hidden_act]91 92 def forward(self, hidden_state):93 return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))94 95 96class Qwen2_5_VisionPatchEmbed(nn.Module):97 def __init__(98 self,99 patch_size: int = 14,100 temporal_patch_size: int = 2,101 in_channels: int = 3,102 embed_dim: int = 1152,103 ) -> None:104 super().__init__()105 self.patch_size = patch_size106 self.temporal_patch_size = temporal_patch_size107 self.in_channels = in_channels108 self.embed_dim = embed_dim109 110 kernel_size = [temporal_patch_size, patch_size, patch_size]111 self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)112 113 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:114 target_dtype = self.proj.weight.dtype115 hidden_states = hidden_states.view(116 -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size117 )118 hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)119 return hidden_states120 121 122class Qwen2_5_VisionRotaryEmbedding(nn.Module):123 inv_freq: torch.Tensor # fix linting for `register_buffer`124 125 def __init__(self, dim: int, theta: float = 10000.0) -> None:126 super().__init__()127 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))128 self.register_buffer("inv_freq", inv_freq, persistent=False)129 130 def forward(self, seqlen: int) -> torch.Tensor:131 seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)132 freqs = torch.outer(seq, self.inv_freq)133 return freqs134 135 136class Qwen2_5_VLPatchMerger(nn.Module):137 def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:138 super().__init__()139 self.hidden_size = context_dim * (spatial_merge_size**2)140 self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6)141 self.mlp = nn.Sequential(142 nn.Linear(self.hidden_size, self.hidden_size),143 nn.GELU(),144 nn.Linear(self.hidden_size, dim),145 )146 147 def forward(self, x: torch.Tensor) -> torch.Tensor:148 x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))149 return x150 151 152def rotate_half(x):153 """Rotates half the hidden dims of the input."""154 x1 = x[..., : x.shape[-1] // 2]155 x2 = x[..., x.shape[-1] // 2 :]156 return torch.cat((-x2, x1), dim=-1)157 158 159def apply_rotary_pos_emb_vision(160 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor161) -> tuple[torch.Tensor, torch.Tensor]:162 orig_q_dtype = q.dtype163 orig_k_dtype = k.dtype164 q, k = q.float(), k.float()165 cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()166 q_embed = (q * cos) + (rotate_half(q) * sin)167 k_embed = (k * cos) + (rotate_half(k) * sin)168 q_embed = q_embed.to(orig_q_dtype)169 k_embed = k_embed.to(orig_k_dtype)170 return q_embed, k_embed171 172 173def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:174 """175 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,176 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)177 """178 batch, num_key_value_heads, slen, head_dim = hidden_states.shape179 if n_rep == 1:180 return hidden_states181 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)182 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)183 184 185def eager_attention_forward(186 module: nn.Module,187 query: torch.Tensor,188 key: torch.Tensor,189 value: torch.Tensor,190 attention_mask: Optional[torch.Tensor],191 scaling: float,192 dropout: float = 0.0,193 **kwargs,194):195 key_states = repeat_kv(key, module.num_key_value_groups)196 value_states = repeat_kv(value, module.num_key_value_groups)197 198 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling199 if attention_mask is not None:200 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]201 attn_weights = attn_weights + causal_mask202 203 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)204 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)205 attn_output = torch.matmul(attn_weights, value_states)206 attn_output = attn_output.transpose(1, 2).contiguous()207 208 return attn_output, attn_weights209 210 211class Qwen2_5_VLVisionAttention(nn.Module):212 def __init__(self, config: Qwen2_5_VLVisionConfig) -> None:213 super().__init__()214 self.dim = config.hidden_size215 self.num_heads = config.num_heads216 self.head_dim = self.dim // self.num_heads217 self.num_key_value_groups = 1 # needed for eager attention218 self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)219 self.proj = nn.Linear(self.dim, self.dim)220 self.scaling = self.head_dim**-0.5221 self.config = config222 self.attention_dropout = 0.0223 self.is_causal = False224 225 def forward(226 self,227 hidden_states: torch.Tensor,228 cu_seqlens: torch.Tensor,229 rotary_pos_emb: Optional[torch.Tensor] = None,230 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,231 **kwargs,232 ) -> torch.Tensor:233 seq_length = hidden_states.shape[0]234 query_states, key_states, value_states = (235 self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)236 )237 cos, sin = position_embeddings238 query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)239 240 query_states = query_states.transpose(0, 1).unsqueeze(0)241 key_states = key_states.transpose(0, 1).unsqueeze(0)242 value_states = value_states.transpose(0, 1).unsqueeze(0)243 244 attention_interface: Callable = eager_attention_forward245 if self.config._attn_implementation != "eager":246 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]247 248 if self.config._attn_implementation == "flash_attention_2":249 # Flash Attention 2: Use cu_seqlens for variable length attention250 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()251 attn_output, _ = attention_interface(252 self,253 query_states,254 key_states,255 value_states,256 attention_mask=None,257 scaling=self.scaling,258 dropout=0.0 if not self.training else self.attention_dropout,259 cu_seq_lens_q=cu_seqlens,260 cu_seq_lens_k=cu_seqlens,261 max_length_q=max_seqlen,262 max_length_k=max_seqlen,263 is_causal=False,264 **kwargs,265 )266 else:267 # Other implementations: Process each chunk separately268 lengths = cu_seqlens[1:] - cu_seqlens[:-1]269 splits = [270 torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)271 ]272 273 attn_outputs = [274 attention_interface(275 self,276 q,277 k,278 v,279 attention_mask=None,280 scaling=self.scaling,281 dropout=0.0 if not self.training else self.attention_dropout,282 is_causal=False,283 **kwargs,284 )[0]285 for q, k, v in zip(*splits)286 ]287 attn_output = torch.cat(attn_outputs, dim=1)288 289 attn_output = attn_output.reshape(seq_length, -1).contiguous()290 attn_output = self.proj(attn_output)291 return attn_output292 293 294class Qwen2_5_VLVisionBlock(GradientCheckpointingLayer):295 def __init__(self, config, attn_implementation: str = "sdpa") -> None:296 super().__init__()297 self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)298 self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6)299 self.attn = Qwen2_5_VLVisionAttention(config=config)300 self.mlp = Qwen2_5_VLMLP(config, bias=True)301 302 def forward(303 self,304 hidden_states: torch.Tensor,305 cu_seqlens: torch.Tensor,306 rotary_pos_emb: Optional[torch.Tensor] = None,307 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,308 **kwargs,309 ) -> torch.Tensor:310 hidden_states = hidden_states + self.attn(311 self.norm1(hidden_states),312 cu_seqlens=cu_seqlens,313 rotary_pos_emb=rotary_pos_emb,314 position_embeddings=position_embeddings,315 **kwargs,316 )317 hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))318 return hidden_states319 320 321@auto_docstring322class Qwen2_5_VLPreTrainedModel(PreTrainedModel):323 config: Qwen2_5_VLConfig324 base_model_prefix = "model"325 supports_gradient_checkpointing = True326 _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]327 _skip_keys_device_placement = "past_key_values"328 _supports_flash_attn = True329 _supports_sdpa = True330 331 _can_compile_fullgraph = True332 _supports_attention_backend = True333 334 335class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):336 config: Qwen2_5_VLVisionConfig337 _no_split_modules = ["Qwen2_5_VLVisionBlock"]338 339 def __init__(self, config, *inputs, **kwargs) -> None:340 super().__init__(config, *inputs, **kwargs)341 self.spatial_merge_size = config.spatial_merge_size342 self.patch_size = config.patch_size343 self.fullatt_block_indexes = config.fullatt_block_indexes344 self.window_size = config.window_size345 self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size346 347 self.patch_embed = Qwen2_5_VisionPatchEmbed(348 patch_size=config.patch_size,349 temporal_patch_size=config.temporal_patch_size,350 in_channels=config.in_channels,351 embed_dim=config.hidden_size,352 )353 354 head_dim = config.hidden_size // config.num_heads355 self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)356 357 self.blocks = nn.ModuleList([Qwen2_5_VLVisionBlock(config) for _ in range(config.depth)])358 self.merger = Qwen2_5_VLPatchMerger(359 dim=config.out_hidden_size,360 context_dim=config.hidden_size,361 spatial_merge_size=config.spatial_merge_size,362 )363 self.gradient_checkpointing = False364 365 def rot_pos_emb(self, grid_thw):366 pos_ids = []367 for t, h, w in grid_thw:368 hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)369 hpos_ids = hpos_ids.reshape(370 h // self.spatial_merge_size,371 self.spatial_merge_size,372 w // self.spatial_merge_size,373 self.spatial_merge_size,374 )375 hpos_ids = hpos_ids.permute(0, 2, 1, 3)376 hpos_ids = hpos_ids.flatten()377 378 wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)379 wpos_ids = wpos_ids.reshape(380 h // self.spatial_merge_size,381 self.spatial_merge_size,382 w // self.spatial_merge_size,383 self.spatial_merge_size,384 )385 wpos_ids = wpos_ids.permute(0, 2, 1, 3)386 wpos_ids = wpos_ids.flatten()387 pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))388 pos_ids = torch.cat(pos_ids, dim=0)389 max_grid_size = grid_thw[:, 1:].max()390 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)391 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)392 return rotary_pos_emb393 394 def get_window_index(self, grid_thw):395 window_index: list = []396 cu_window_seqlens: list = [0]397 window_index_id = 0398 vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size399 400 for grid_t, grid_h, grid_w in grid_thw:401 llm_grid_h, llm_grid_w = (402 grid_h // self.spatial_merge_size,403 grid_w // self.spatial_merge_size,404 )405 index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)406 pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size407 pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size408 num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size409 num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size410 index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)411 index_padded = index_padded.reshape(412 grid_t,413 num_windows_h,414 vit_merger_window_size,415 num_windows_w,416 vit_merger_window_size,417 )418 index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(419 grid_t,420 num_windows_h * num_windows_w,421 vit_merger_window_size,422 vit_merger_window_size,423 )424 seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)425 index_padded = index_padded.reshape(-1)426 index_new = index_padded[index_padded != -100]427 window_index.append(index_new + window_index_id)428 cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]429 cu_window_seqlens.extend(cu_seqlens_tmp.tolist())430 window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()431 window_index = torch.cat(window_index, dim=0)432 433 return window_index, cu_window_seqlens434 435 def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:436 """437 Args:438 hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):439 The final hidden states of the model.440 grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):441 The temporal, height and width of feature shape of each image in LLM.442 443 Returns:444 `torch.Tensor`: hidden_states.445 """446 hidden_states = self.patch_embed(hidden_states)447 rotary_pos_emb = self.rot_pos_emb(grid_thw)448 window_index, cu_window_seqlens = self.get_window_index(grid_thw)449 cu_window_seqlens = torch.tensor(450 cu_window_seqlens,451 device=hidden_states.device,452 dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,453 )454 cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)455 456 seq_len, _ = hidden_states.size()457 hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)458 hidden_states = hidden_states[window_index, :, :]459 hidden_states = hidden_states.reshape(seq_len, -1)460 rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)461 rotary_pos_emb = rotary_pos_emb[window_index, :, :]462 rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)463 emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)464 position_embeddings = (emb.cos(), emb.sin())465 466 cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(467 dim=0,468 # Select dtype based on the following factors:469 # - FA2 requires that cu_seqlens_q must have dtype int32470 # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw471 # See https://github.com/huggingface/transformers/pull/34852 for more information472 dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,473 )474 cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)475 476 for layer_num, blk in enumerate(self.blocks):477 if layer_num in self.fullatt_block_indexes:478 cu_seqlens_now = cu_seqlens479 else:480 cu_seqlens_now = cu_window_seqlens481 482 hidden_states = blk(483 hidden_states,484 cu_seqlens=cu_seqlens_now,485 position_embeddings=position_embeddings,486 **kwargs,487 )488 489 hidden_states = self.merger(hidden_states)490 reverse_indices = torch.argsort(window_index)491 hidden_states = hidden_states[reverse_indices, :]492 493 return hidden_states494 495 496@dataclass497@auto_docstring(498 custom_intro="""499 Base class for Llava outputs, with hidden states and attentions.500 """501)502class Qwen2_5_VLModelOutputWithPast(ModelOutput):503 r"""504 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):505 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).506 507 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see508 `past_key_values` input) to speed up sequential decoding.509 rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):510 The rope index difference between sequence length and multimodal rope.511 """512 513 last_hidden_state: Optional[torch.FloatTensor] = None514 past_key_values: Optional[Cache] = None515 hidden_states: Optional[tuple[torch.FloatTensor]] = None516 attentions: Optional[tuple[torch.FloatTensor]] = None517 rope_deltas: Optional[torch.LongTensor] = None518 519 520class Qwen2_5_VLRotaryEmbedding(nn.Module):521 inv_freq: torch.Tensor # fix linting for `register_buffer`522 523 def __init__(self, config: Qwen2_5_VLTextConfig, device=None):524 super().__init__()525 # BC: "rope_type" was originally "type"526 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:527 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))528 else:529 self.rope_type = "default"530 self.max_seq_len_cached = config.max_position_embeddings531 self.original_max_seq_len = config.max_position_embeddings532 533 self.config = config534 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]535 536 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)537 self.register_buffer("inv_freq", inv_freq, persistent=False)538 self.original_inv_freq = self.inv_freq539 540 @torch.no_grad()541 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)542 def forward(self, x, position_ids):543 # In contrast to other models, Qwen2_5_VL has different position ids for the grids544 # So we expand the inv_freq to shape (3, ...)545 inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)546 position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)547 548 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"549 with torch.autocast(device_type=device_type, enabled=False): # Force float32550 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)551 emb = torch.cat((freqs, freqs), dim=-1)552 cos = emb.cos() * self.attention_scaling553 sin = emb.sin() * self.attention_scaling554 555 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)556 557 558class Qwen2MLP(nn.Module):559 def __init__(self, config):560 super().__init__()561 self.config = config562 self.hidden_size = config.hidden_size563 self.intermediate_size = config.intermediate_size564 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)565 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)566 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)567 self.act_fn = ACT2FN[config.hidden_act]568 569 def forward(self, x):570 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))571 return down_proj572 573 574def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):575 """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).576 577 Explanation:578 Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding579 sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For580 vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.581 Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.582 For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,583 height and width) of text embedding is always the same, so the text embedding rotary position embedding has no584 difference with modern LLMs.585 586 Args:587 q (`torch.Tensor`): The query tensor.588 k (`torch.Tensor`): The key tensor.589 cos (`torch.Tensor`): The cosine part of the rotary embedding.590 sin (`torch.Tensor`): The sine part of the rotary embedding.591 position_ids (`torch.Tensor`):592 The position indices of the tokens corresponding to the query and key tensors. For example, this can be593 used to pass offsetted position ids when working with a KV-cache.594 mrope_section(`List(int)`):595 Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.596 unsqueeze_dim (`int`, *optional*, defaults to 1):597 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and598 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note599 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and600 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes601 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have602 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.603 Returns:604 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.605 """606 mrope_section = mrope_section * 2607 cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(608 unsqueeze_dim609 )610 sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(611 unsqueeze_dim612 )613 614 q_embed = (q * cos) + (rotate_half(q) * sin)615 k_embed = (k * cos) + (rotate_half(k) * sin)616 return q_embed, k_embed617 618 619class Qwen2_5_VLAttention(nn.Module):620 """621 Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer622 and "Generating Long Sequences with Sparse Transformers".623 """624 625 def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: Optional[int] = None):626 super().__init__()627 self.config = config628 self.layer_idx = layer_idx629 if layer_idx is None:630 logger.warning_once(631 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "632 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "633 "when creating this class."634 )635 636 self.hidden_size = config.hidden_size637 self.num_heads = config.num_attention_heads638 self.head_dim = self.hidden_size // self.num_heads639 self.num_key_value_heads = config.num_key_value_heads640 self.num_key_value_groups = self.num_heads // self.num_key_value_heads641 self.is_causal = True642 self.attention_dropout = config.attention_dropout643 self.rope_scaling = config.rope_scaling644 self.scaling = self.head_dim**-0.5645 646 if (self.head_dim * self.num_heads) != self.hidden_size:647 raise ValueError(648 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"649 f" and `num_heads`: {self.num_heads})."650 )651 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)652 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)653 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)654 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)655 self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None656 657 self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config)658 659 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")660 def forward(661 self,662 hidden_states: torch.Tensor,663 attention_mask: Optional[torch.Tensor] = None,664 position_ids: Optional[torch.LongTensor] = None,665 past_key_values: Optional[Cache] = None,666 output_attentions: bool = False,667 use_cache: bool = False,668 cache_position: Optional[torch.LongTensor] = None,669 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC670 **kwargs: Unpack[FlashAttentionKwargs],671 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:672 bsz, q_len, _ = hidden_states.size()673 674 query_states = self.q_proj(hidden_states)675 key_states = self.k_proj(hidden_states)676 value_states = self.v_proj(hidden_states)677 678 query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)679 key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)680 value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)681 682 cos, sin = position_embeddings683 query_states, key_states = apply_multimodal_rotary_pos_emb(684 query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]685 )686 687 if past_key_values is not None:688 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models689 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)690 691 attention_interface: Callable = eager_attention_forward692 if self.config._attn_implementation != "eager":693 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]694 695 attn_output, attn_weights = attention_interface(696 self,697 query_states,698 key_states,699 value_states,700 attention_mask,701 dropout=0.0 if not self.training else self.attention_dropout,702 scaling=self.scaling,703 sliding_window=self.sliding_window,704 position_ids=position_ids, # pass positions for FA2705 **kwargs,706 )707 708 attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()709 attn_output = self.o_proj(attn_output)710 return attn_output, attn_weights711 712 713class Qwen2_5_VLDecoderLayer(GradientCheckpointingLayer):714 def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int):715 super().__init__()716 self.hidden_size = config.hidden_size717 718 if config.use_sliding_window and config._attn_implementation != "flash_attention_2":719 logger.warning_once(720 f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "721 "unexpected results may be encountered."722 )723 self.self_attn = Qwen2_5_VLAttention(config, layer_idx)724 725 self.mlp = Qwen2MLP(config)726 self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)727 self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)728 self.attention_type = config.layer_types[layer_idx]729 730 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")731 def forward(732 self,733 hidden_states: torch.Tensor,734 attention_mask: Optional[torch.Tensor] = None,735 position_ids: Optional[torch.LongTensor] = None,736 past_key_values: Optional[Cache] = None,737 output_attentions: Optional[bool] = False,738 use_cache: Optional[bool] = False,739 cache_position: Optional[torch.LongTensor] = None,740 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC741 **kwargs: Unpack[FlashAttentionKwargs],742 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:743 """744 Args:745 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`746 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size747 `(batch, sequence_length)` where padding elements are indicated by 0.748 output_attentions (`bool`, *optional*):749 Whether or not to return the attentions tensors of all attention layers. See `attentions` under750 returned tensors for more detail.751 use_cache (`bool`, *optional*):752 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding753 (see `past_key_values`).754 past_key_values (`Cache`, *optional*): cached past key and value projection states755 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):756 Indices depicting the position of the input sequence tokens in the sequence.757 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):758 Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,759 with `head_dim` being the embedding dimension of each attention head.760 kwargs (`dict`, *optional*):761 Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code762 into the model763 """764 765 residual = hidden_states766 767 hidden_states = self.input_layernorm(hidden_states)768 769 # Self Attention770 hidden_states, self_attn_weights = self.self_attn(771 hidden_states=hidden_states,772 attention_mask=attention_mask,773 position_ids=position_ids,774 past_key_values=past_key_values,775 output_attentions=output_attentions,776 use_cache=use_cache,777 cache_position=cache_position,778 position_embeddings=position_embeddings,779 **kwargs,780 )781 hidden_states = residual + hidden_states782 783 # Fully Connected784 residual = hidden_states785 hidden_states = self.post_attention_layernorm(hidden_states)786 hidden_states = self.mlp(hidden_states)787 hidden_states = residual + hidden_states788 789 outputs = (hidden_states,)790 791 if output_attentions:792 outputs += (self_attn_weights,)793 794 return outputs795 796 797@auto_docstring798class Qwen2_5_VLTextModel(Qwen2_5_VLPreTrainedModel):799 config: Qwen2_5_VLTextConfig800 801 def __init__(self, config: Qwen2_5_VLTextConfig):802 super().__init__(config)803 self.padding_idx = config.pad_token_id804 self.vocab_size = config.vocab_size805 806 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)807 self.layers = nn.ModuleList(808 [Qwen2_5_VLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]809 )810 self._attn_implementation = config._attn_implementation811 self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)812 self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config)813 self.has_sliding_layers = "sliding_attention" in self.config.layer_types814 815 self.gradient_checkpointing = False816 # Initialize weights and apply final processing817 self.post_init()818 819 @auto_docstring820 def forward(821 self,822 input_ids: Optional[torch.LongTensor] = None,823 attention_mask: Optional[torch.Tensor] = None,824 position_ids: Optional[torch.LongTensor] = None,825 past_key_values: Optional[Cache] = None,826 inputs_embeds: Optional[torch.FloatTensor] = None,827 use_cache: Optional[bool] = None,828 output_attentions: Optional[bool] = None,829 output_hidden_states: Optional[bool] = None,830 return_dict: Optional[bool] = None,831 cache_position: Optional[torch.LongTensor] = None,832 **kwargs: Unpack[FlashAttentionKwargs],833 ) -> Union[tuple, BaseModelOutputWithPast]:834 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions835 output_hidden_states = (836 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states837 )838 use_cache = use_cache if use_cache is not None else self.config.use_cache839 840 return_dict = return_dict if return_dict is not None else self.config.use_return_dict841 842 if (input_ids is None) ^ (inputs_embeds is not None):843 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")844 845 if self.gradient_checkpointing and self.training:846 if use_cache:847 logger.warning_once(848 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."849 )850 use_cache = False851 852 # torch.jit.trace() doesn't support cache objects in the output853 if use_cache and past_key_values is None and not torch.jit.is_tracing():854 past_key_values = DynamicCache(config=self.config)855 856 if inputs_embeds is None:857 inputs_embeds = self.embed_tokens(input_ids)858 859 if cache_position is None:860 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0861 cache_position = torch.arange(862 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device863 )864 865 # the hard coded `3` is for temporal, height and width.866 if position_ids is None:867 position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)868 elif position_ids.ndim == 2:869 position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)870 871 # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions872 # where each dim indicates visual spatial positions for temporal/height/width grids.873 # There are two scenarios when FA2-like packed masking might be activated.874 # 1. User specifically passed packed `position_ids` and no attention mask.875 # In this case we expect the useer to create correct position ids for all 3 grids876 # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]877 # 2. User runs forward with no attention mask and no position ids. In this case, position ids878 # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are879 # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass880 # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`881 if position_ids.ndim == 3 and position_ids.shape[0] == 4:882 text_position_ids = position_ids[0]883 position_ids = position_ids[1:]884 else:885 # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids886 text_position_ids = None887 888 # It may already have been prepared by e.g. `generate`889 if not isinstance(causal_mask_mapping := attention_mask, dict):890 # Prepare mask arguments891 mask_kwargs = {892 "config": self.config,893 "input_embeds": inputs_embeds,894 "attention_mask": attention_mask,895 "cache_position": cache_position,896 "past_key_values": past_key_values,897 "position_ids": text_position_ids,898 }899 # Create the masks900 causal_mask_mapping = {901 "full_attention": create_causal_mask(**mask_kwargs),902 }903 # The sliding window alternating layers are not always activated depending on the config904 if self.has_sliding_layers:905 causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)906 907 hidden_states = inputs_embeds908 909 # create position embeddings to be shared across the decoder layers910 position_embeddings = self.rotary_emb(hidden_states, position_ids)911 912 # decoder layers913 all_hidden_states = () if output_hidden_states else None914 all_self_attns = () if output_attentions else None915 916 for decoder_layer in self.layers:917 if output_hidden_states:918 all_hidden_states += (hidden_states,)919 920 layer_outputs = decoder_layer(921 hidden_states,922 attention_mask=causal_mask_mapping[decoder_layer.attention_type],923 position_ids=text_position_ids,924 past_key_values=past_key_values,925 output_attentions=output_attentions,926 use_cache=use_cache,927 cache_position=cache_position,928 position_embeddings=position_embeddings,929 **kwargs,930 )931 932 hidden_states = layer_outputs[0]933 934 if output_attentions:935 all_self_attns += (layer_outputs[1],)936 937 hidden_states = self.norm(hidden_states)938 939 # add hidden states from the last decoder layer940 if output_hidden_states:941 all_hidden_states += (hidden_states,)942 943 if not return_dict:944 return tuple(945 v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None946 )947 return BaseModelOutputWithPast(948 last_hidden_state=hidden_states,949 past_key_values=past_key_values,950 hidden_states=all_hidden_states,951 attentions=all_self_attns,952 )953 954 955@auto_docstring956class Qwen2_5_VLModel(Qwen2_5_VLPreTrainedModel):957 base_model_prefix = ""958 _checkpoint_conversion_mapping = {"^model": "language_model"}959 # Reference: fix gemma3 grad acc #37208960 accepts_loss_kwargs = False961 config: Qwen2_5_VLConfig962 _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]963 964 def __init__(self, config):965 super().__init__(config)966 self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)967 self.language_model = Qwen2_5_VLTextModel._from_config(config.text_config)968 self.rope_deltas = None # cache rope_deltas here969 # Initialize weights and apply final processing970 self.post_init()971 972 def get_input_embeddings(self):973 return self.language_model.get_input_embeddings()974 975 def set_input_embeddings(self, value):976 self.language_model.set_input_embeddings(value)977 978 def set_decoder(self, decoder):979 self.language_model = decoder980 981 def get_decoder(self):982 return self.language_model983 984 def get_rope_index(985 self,986 input_ids: Optional[torch.LongTensor] = None,987 image_grid_thw: Optional[torch.LongTensor] = None,988 video_grid_thw: Optional[torch.LongTensor] = None,989 second_per_grid_ts: Optional[torch.Tensor] = None,990 attention_mask: Optional[torch.Tensor] = None,991 ) -> tuple[torch.Tensor, torch.Tensor]:992 """993 Calculate the 3D rope index based on image and video's temporal, height and width in LLM.994 995 Explanation:996 Each embedding sequence contains vision embedding and text embedding or just contains text embedding.997 998 For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.999 Examples:1000 input_ids: [T T T T T], here T is for text.1001 temporal position_ids: [0, 1, 2, 3, 4]1002 height position_ids: [0, 1, 2, 3, 4]1003 width position_ids: [0, 1, 2, 3, 4]1004 1005 For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part1006 and 1D rotary position embedding for text part.1007 Examples:1008 Temporal (Time): 3 patches, representing different segments of the video in time.1009 Height: 2 patches, dividing each frame vertically.1010 Width: 2 patches, dividing each frame horizontally.1011 We also have some important parameters:1012 fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.1013 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.1014 temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.1015 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.1016 input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.1017 vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]1018 vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]1019 vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]1020 text temporal position_ids: [101, 102, 103, 104, 105]1021 text height position_ids: [101, 102, 103, 104, 105]1022 text width position_ids: [101, 102, 103, 104, 105]1023 Here we calculate the text start position_ids as the max vision position_ids plus 1.1024 1025 Args:1026 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):1027 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide1028 it.1029 image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):1030 The temporal, height and width of feature shape of each image in LLM.1031 video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):1032 The temporal, height and width of feature shape of each video in LLM.1033 second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):1034 The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.1035 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1036 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:1037 1038 - 1 for tokens that are **not masked**,1039 - 0 for tokens that are **masked**.1040 1041 Returns:1042 position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)1043 mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)1044 """1045 spatial_merge_size = self.config.vision_config.spatial_merge_size1046 image_token_id = self.config.image_token_id1047 video_token_id = self.config.video_token_id1048 vision_start_token_id = self.config.vision_start_token_id1049 mrope_position_deltas = []1050 if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):1051 total_input_ids = input_ids1052 if attention_mask is not None:1053 attention_mask = attention_mask == 11054 position_ids = torch.ones(1055 3,1056 input_ids.shape[0],1057 input_ids.shape[1],1058 dtype=input_ids.dtype,1059 device=input_ids.device,1060 )1061 image_index, video_index = 0, 01062 for i, input_ids in enumerate(total_input_ids):1063 if attention_mask is not None:1064 input_ids = input_ids[attention_mask[i]]1065 image_nums, video_nums = 0, 01066 vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)1067 vision_tokens = input_ids[vision_start_indices + 1]1068 image_nums = (vision_tokens == image_token_id).sum()1069 video_nums = (vision_tokens == video_token_id).sum()1070 input_tokens = input_ids.tolist()1071 llm_pos_ids_list: list = []1072 st = 01073 remain_images, remain_videos = image_nums, video_nums1074 for _ in range(image_nums + video_nums):1075 if image_token_id in input_tokens and remain_images > 0:1076 ed_image = input_tokens.index(image_token_id, st)1077 else:1078 ed_image = len(input_tokens) + 11079 if video_token_id in input_tokens and remain_videos > 0:1080 ed_video = input_tokens.index(video_token_id, st)1081 else:1082 ed_video = len(input_tokens) + 11083 if ed_image < ed_video:1084 t, h, w = (1085 image_grid_thw[image_index][0],1086 image_grid_thw[image_index][1],1087 image_grid_thw[image_index][2],1088 )1089 second_per_grid_t = 01090 image_index += 11091 remain_images -= 11092 ed = ed_image1093 1094 else:1095 t, h, w = (1096 video_grid_thw[video_index][0],1097 video_grid_thw[video_index][1],1098 video_grid_thw[video_index][2],1099 )1100 if second_per_grid_ts is not None:1101 second_per_grid_t = second_per_grid_ts[video_index]1102 else:1103 second_per_grid_t = 1.01104 video_index += 11105 remain_videos -= 11106 ed = ed_video1107 llm_grid_t, llm_grid_h, llm_grid_w = (1108 t.item(),1109 h.item() // spatial_merge_size,1110 w.item() // spatial_merge_size,1111 )1112 text_len = ed - st1113 1114 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 01115 llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)1116 1117 range_tensor = torch.arange(llm_grid_t).view(-1, 1)1118 expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)1119 1120 ## normalize type, send to device.1121 second_per_grid_t = torch.as_tensor(1122 second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device1123 )1124 1125 time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second1126 1127 time_tensor_long = time_tensor.long()1128 t_index = time_tensor_long.flatten()1129 1130 h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()1131 w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()1132 llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)1133 st = ed + llm_grid_t * llm_grid_h * llm_grid_w1134 1135 if st < len(input_tokens):1136 st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 01137 text_len = len(input_tokens) - st1138 llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)1139 1140 llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)1141 if attention_mask is not None:1142 position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device)1143 else:1144 position_ids[..., i, :] = llm_positions.to(position_ids.device)1145 mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))1146 mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device)1147 return position_ids, mrope_position_deltas1148 else:1149 if attention_mask is not None:1150 position_ids = attention_mask.long().cumsum(-1) - 11151 position_ids.masked_fill_(attention_mask == 0, 1)1152 position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)1153 max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]1154 mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]1155 else:1156 position_ids = (1157 torch.arange(input_ids.shape[1], device=input_ids.device)1158 .view(1, 1, -1)1159 .expand(3, input_ids.shape[0], -1)1160 )1161 mrope_position_deltas = torch.zeros(1162 [input_ids.shape[0], 1],1163 device=input_ids.device,1164 dtype=input_ids.dtype,1165 )1166 1167 return position_ids, mrope_position_deltas1168 1169 def get_video_features(1170 self, pixel_values_videos: torch.FloatTensor, video_grid_thw: Optional[torch.LongTensor] = None1171 ):1172 """1173 Encodes videos into continuous embeddings that can be forwarded to the language model.1174 1175 Args:1176 pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):1177 The tensors corresponding to the input videos.1178 video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):1179 The temporal, height and width of feature shape of each video in LLM.1180 """1181 pixel_values_videos = pixel_values_videos.type(self.visual.dtype)1182 video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)1183 split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()1184 video_embeds = torch.split(video_embeds, split_sizes)1185 return video_embeds1186 1187 def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):1188 """1189 Encodes images into continuous embeddings that can be forwarded to the language model.1190 1191 Args:1192 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):1193 The tensors corresponding to the input images.1194 image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):1195 The temporal, height and width of feature shape of each image in LLM.1196 """1197 pixel_values = pixel_values.type(self.visual.dtype)1198 image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)1199 split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()1200 image_embeds = torch.split(image_embeds, split_sizes)