diff --git a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py index a0198139bc02..c35b214cdae2 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py @@ -1106,52 +1106,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, - mm_token_type_ids: torch.IntTensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1161,8 +1206,6 @@ def get_rope_index( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`, *optional*): - Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). Returns: position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) @@ -1173,64 +1216,55 @@ def get_rope_index( spatial_merge_size = self.config.vision_config.spatial_merge_size mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - # If we don't have `mm_token_type_ids`, then we have text tokens only (== 0) - if mm_token_type_ids is None: - input_token_type = torch.zeros_like(input_ids)[attention_mask[i] == 1].tolist() - else: - input_token_type = mm_token_type_ids[i, attention_mask[i] == 1].tolist() + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 llm_pos_ids_list = [] for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 # text == 0 if modality_type == 0: text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len # image == 1, video == 2 else: - grid_thw = image_grid_thw if modality_type == 1 else video_grid_thw - mm_index = image_index if modality_type == 1 else video_index + grid_thw = next(grid_iters[modality_type]) t_merge_size = 1 if modality_type == 1 else temporal_merge_size - t, h, w = ( - grid_thw[mm_index][0], - grid_thw[mm_index][1], - grid_thw[mm_index][2], - ) - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item() // t_merge_size, - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, t_merge_size, spatial_merge_size, device=input_ids.device ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - if modality_type == 1: - image_index += 1 - else: - video_index += 1 + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1329,10 +1363,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, - mm_token_type_ids: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1354,7 +1392,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = (position_ids + delta).to(device=inputs_embeds.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1728,8 +1766,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 68a181eccfa3..bac6b5098cda 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -78,14 +78,13 @@ from ..qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionPatchEmbed, Qwen2_5_VisionRotaryEmbedding, - Qwen2_5_VLModel, Qwen2_5_VLPreTrainedModel, Qwen2_5_VLVisionAttention, Qwen2_5_VLVisionBlock, ) from ..qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig from ..qwen2_vl.image_processing_qwen2_vl import smart_resize -from ..qwen2_vl.modeling_qwen2_vl import Qwen2VisionTransformerPretrainedModel, VisionMlp +from ..qwen2_vl.modeling_qwen2_vl import Qwen2VisionTransformerPretrainedModel, Qwen2VLModel, VisionMlp logger = logging.get_logger(__name__) @@ -1038,8 +1037,10 @@ def forward(self, hidden_states, grid_thw): return hidden_states -class Ernie4_5_VL_MoeModel(Qwen2_5_VLModel): +class Ernie4_5_VL_MoeModel(Qwen2VLModel): _checkpoint_conversion_mapping = {"^norm": "language_model.norm"} + config: Ernie4_5_VL_MoeConfig + _no_split_modules = ["Ernie4_5_VL_MoeDecoderLayer", "Ernie4_5_VL_MoeVisionBlock"] def __init__(self, config: Ernie4_5_VL_MoeConfig): super().__init__(config) @@ -1050,50 +1051,39 @@ def __init__(self, config: Ernie4_5_VL_MoeConfig): def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, - mm_token_type_ids: torch.IntTensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1103,8 +1093,6 @@ def get_rope_index( - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. - mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`, *optional*): - Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). Returns: position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) @@ -1115,64 +1103,55 @@ def get_rope_index( spatial_merge_size = self.config.vision_config.spatial_merge_size mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - # If we don't have `mm_token_type_ids`, then we have text tokens only (== 0) - if mm_token_type_ids is None: - input_token_type = torch.zeros_like(input_ids)[attention_mask[i] == 1].tolist() - else: - input_token_type = mm_token_type_ids[i, attention_mask[i] == 1].tolist() + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 llm_pos_ids_list = [] for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 # text == 0 if modality_type == 0: text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len # image == 1, video == 2 else: - grid_thw = image_grid_thw if modality_type == 1 else video_grid_thw - mm_index = image_index if modality_type == 1 else video_index + grid_thw = next(grid_iters[modality_type]) t_merge_size = 1 if modality_type == 1 else temporal_merge_size - t, h, w = ( - grid_thw[mm_index][0], - grid_thw[mm_index][1], - grid_thw[mm_index][2], + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, t_merge_size, spatial_merge_size, device=input_ids.device ) - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item() // t_merge_size, - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - if modality_type == 1: - image_index += 1 - else: - video_index += 1 + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1210,45 +1189,6 @@ def get_image_features( image_outputs.pooler_output = image_embeds return image_outputs - def compute_3d_position_ids( - self, - input_ids: torch.Tensor | None, - inputs_embeds: torch.Tensor | None, - image_grid_thw: torch.Tensor | None = None, - video_grid_thw: torch.Tensor | None = None, - attention_mask: torch.Tensor | None = None, - past_key_values: torch.Tensor | None = None, - mm_token_type_ids: torch.Tensor | None = None, - ) -> torch.Tensor | None: - past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) - - if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): - position_ids, rope_deltas = self.get_rope_index( - input_ids, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - attention_mask=attention_mask, - mm_token_type_ids=mm_token_type_ids, - ) - self.rope_deltas = rope_deltas - # Use pre-calculated rope-deltas to infer correct 3D position ids - elif self.rope_deltas is not None: - batch_size, seq_length, _ = inputs_embeds.shape - if attention_mask is not None: - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids = position_ids.masked_fill(attention_mask == 0, 0) - position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device) - else: - position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) - position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) - delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = (position_ids + delta).to(device=inputs_embeds.device) - else: - # Can't build correct 3D positions. Let the model infer it from `cache_position` - position_ids = None - return position_ids - @auto_docstring @can_return_tuple def forward( diff --git a/src/transformers/models/glm46v/modeling_glm46v.py b/src/transformers/models/glm46v/modeling_glm46v.py index 3dbf6b8adf22..f84278963268 100644 --- a/src/transformers/models/glm46v/modeling_glm46v.py +++ b/src/transformers/models/glm46v/modeling_glm46v.py @@ -99,51 +99,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -158,93 +204,71 @@ def get_rope_index( position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_start_token_id = self.config.video_start_token_id - video_end_token_id = self.config.video_end_token_id mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - video_group_index = 0 - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - - input_token_type = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False - if token == image_token_id and not video_check_flg: - input_token_type.append("image") - elif token == image_token_id and video_check_flg: - input_token_type.append("video") - else: - input_token_type.append("text") + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + video_group_index = 0 llm_pos_ids_list = [] - video_frame_num = 1 for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - if modality_type == "image": - t, h, w = image_grid_thw_list[image_index] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - image_index += 1 - video_frame_num = 1 - elif modality_type == "video": - _, h, w = video_grid_thw_list[video_index] - t = video_frame_num - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - video_group_index += 1 - if video_group_index >= video_grid_thw_list[video_index][0]: - video_index += 1 - video_group_index = 0 - video_frame_num += 1 + current_pos += text_len + # image == 1, video == 2 else: - text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - video_frame_num = 1 + # GLM46V splits video into segments per frame but there's only one `grid_thw` + # per whole video. We can't exhaus the iterator and have to re-use the grid + # while processing the same video! + if modality_type == 2: + if video_group_index == 0: + grid_thw = next(grid_iters[modality_type]) + video_group_index += 1 + video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index + else: + grid_thw = next(grid_iters[modality_type]) + + # Videos are processed per frame separately, each temporal grid is always `1` + temp_merge_size = grid_thw[0] + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -351,9 +375,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -361,6 +390,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -374,7 +404,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -394,6 +424,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm46VModelOutputWithPast: @@ -431,6 +462,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -543,6 +575,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -595,6 +628,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -680,8 +714,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/glm46v/processing_glm46v.py b/src/transformers/models/glm46v/processing_glm46v.py index 99bc1cbd8dcd..3b71afd1183b 100644 --- a/src/transformers/models/glm46v/processing_glm46v.py +++ b/src/transformers/models/glm46v/processing_glm46v.py @@ -37,7 +37,7 @@ class Glm46VProcessorKwargs(ProcessingKwargs, total=False): "text_kwargs": { "padding": False, "return_token_type_ids": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -59,6 +59,8 @@ def __init__(self, image_processor=None, tokenizer=None, video_processor=None, c else tokenizer.convert_tokens_to_ids(self.video_token) ) super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) + self.video_start_id = tokenizer.convert_tokens_to_ids("<|begin_of_video|>") + self.video_end_id = tokenizer.convert_tokens_to_ids("<|end_of_video|>") @auto_docstring def __call__( @@ -169,7 +171,16 @@ def __call__( if return_mm_token_type_ids: array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) - mm_token_type_ids[array_ids == self.image_token_id] = 1 + + # Replace 0 -> 2 only inside video segments because Glm46V + # uses the same special token to denote images and video + # Otherwise replace 0 -> 1 for image modality + starts = np.cumsum(array_ids == self.video_start_id, axis=1) + ends = np.cumsum(array_ids == self.video_end_id, axis=1) + is_video_modality = starts > ends + + mm_token_type_ids[(array_ids == self.image_token_id) & is_video_modality] = 2 + mm_token_type_ids[(array_ids == self.image_token_id) & (~is_video_modality)] = 1 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) @@ -238,6 +249,12 @@ def post_process_image_text_to_text( **kwargs, ) + @property + def model_input_names(self): + model_input_names = super().model_input_names + model_input_names.append("mm_token_type_ids") + return model_input_names + def replace_frame_token_id(self, timestamp_sec): return f"<|begin_of_image|>{self.image_token}<|end_of_image|>{timestamp_sec:.1f} seconds" diff --git a/src/transformers/models/glm4v/modeling_glm4v.py b/src/transformers/models/glm4v/modeling_glm4v.py index b5a763c22c81..cc65db0ced9e 100644 --- a/src/transformers/models/glm4v/modeling_glm4v.py +++ b/src/transformers/models/glm4v/modeling_glm4v.py @@ -951,51 +951,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1010,93 +1056,71 @@ def get_rope_index( position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_start_token_id = self.config.video_start_token_id - video_end_token_id = self.config.video_end_token_id mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - video_group_index = 0 - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - - input_token_type = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False - if token == image_token_id and not video_check_flg: - input_token_type.append("image") - elif token == image_token_id and video_check_flg: - input_token_type.append("video") - else: - input_token_type.append("text") + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + video_group_index = 0 llm_pos_ids_list = [] - video_frame_num = 1 for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - if modality_type == "image": - t, h, w = image_grid_thw_list[image_index] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - image_index += 1 - video_frame_num = 1 - elif modality_type == "video": - _, h, w = video_grid_thw_list[video_index] - t = video_frame_num - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - video_group_index += 1 - if video_group_index >= video_grid_thw_list[video_index][0]: - video_index += 1 - video_group_index = 0 - video_frame_num += 1 + current_pos += text_len + # image == 1, video == 2 else: - text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - video_frame_num = 1 + # GLM4V splits video into segments per frame but there's only one `grid_thw` + # per whole video. We can't exhaus the iterator and have to re-use the grid + # while processing the same video! + if modality_type == 2: + if video_group_index == 0: + grid_thw = next(grid_iters[modality_type]) + video_group_index += 1 + video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index + else: + grid_thw = next(grid_iters[modality_type]) + + # Videos are processed per frame separately, each temporal grid is always `1` + temp_merge_size = grid_thw[0] + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1203,9 +1227,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1213,6 +1242,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1226,7 +1256,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1246,6 +1276,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm4vModelOutputWithPast: @@ -1283,6 +1314,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1395,6 +1427,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1447,6 +1480,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1532,8 +1566,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index 31365eb69c8b..3e607eedfc75 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -956,155 +956,6 @@ def __init__(self, config): super().__init__(config) self.visual = Glm4vVisionModel._from_config(config.vision_config) - def get_rope_index( - self, - input_ids: torch.LongTensor | None = None, - image_grid_thw: torch.LongTensor | None = None, - video_grid_thw: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. - - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide - it. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - Returns: - position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) - mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) - """ - - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_start_token_id = self.config.video_start_token_id - video_end_token_id = self.config.video_end_token_id - - mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( - 3, - input_ids.shape[0], - input_ids.shape[1], - dtype=input_ids.dtype, - device=input_ids.device, - ) - image_index, video_index = 0, 0 - video_group_index = 0 - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - - input_token_type = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False - if token == image_token_id and not video_check_flg: - input_token_type.append("image") - elif token == image_token_id and video_check_flg: - input_token_type.append("video") - else: - input_token_type.append("text") - input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): - group = list(group) - start_index = group[0][0] - end_index = group[-1][0] + 1 - input_type_group.append((key, start_index, end_index)) - llm_pos_ids_list = [] - video_frame_num = 1 - for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - if modality_type == "image": - t, h, w = image_grid_thw_list[image_index] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - image_index += 1 - video_frame_num = 1 - elif modality_type == "video": - _, h, w = video_grid_thw_list[video_index] - t = video_frame_num - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - video_group_index += 1 - if video_group_index >= video_grid_thw_list[video_index][0]: - video_index += 1 - video_group_index = 0 - video_frame_num += 1 - else: - text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - video_frame_num = 1 - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) - return position_ids, mrope_position_deltas - @can_return_tuple @auto_docstring def get_video_features( @@ -1178,6 +1029,123 @@ def get_placeholder_mask( ) return special_image_mask, special_video_mask + def get_rope_index( + self, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + + mrope_position_deltas = [] + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + video_group_index = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 + else: + # GLM4V splits video into segments per frame but there's only one `grid_thw` + # per whole video. We can't exhaus the iterator and have to re-use the grid + # while processing the same video! + if modality_type == 2: + if video_group_index == 0: + grid_thw = next(grid_iters[modality_type]) + video_group_index += 1 + video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index + else: + grid_thw = next(grid_iters[modality_type]) + + # Videos are processed per frame separately, each temporal grid is always `1` + temp_merge_size = grid_thw[0] + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + @auto_docstring @can_return_tuple def forward( @@ -1192,6 +1160,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm4vModelOutputWithPast: @@ -1229,6 +1198,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1266,6 +1236,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1318,6 +1289,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1447,7 +1419,7 @@ class Glm4vProcessorKwargs(Qwen2VLProcessorKwargs): "text_kwargs": { "padding": False, "return_token_type_ids": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -1458,6 +1430,8 @@ def __init__(self, image_processor=None, tokenizer=None, video_processor=None, c super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) self.image_token = "<|image|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token self.video_token = "<|video|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token + self.video_start_id = tokenizer.convert_tokens_to_ids("<|begin_of_video|>") + self.video_end_id = tokenizer.convert_tokens_to_ids("<|end_of_video|>") def __call__( self, @@ -1567,7 +1541,16 @@ def __call__( if return_mm_token_type_ids: array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) - mm_token_type_ids[array_ids == self.image_token_id] = 1 + + # Replace 0 -> 2 only inside video segments because GLM4v + # uses the same special token to denote images and video + # Otherwise replace 0 -> 1 for image modality + starts = np.cumsum(array_ids == self.video_start_id, axis=1) + ends = np.cumsum(array_ids == self.video_end_id, axis=1) + is_video_modality = starts > ends + + mm_token_type_ids[(array_ids == self.image_token_id) & is_video_modality] = 2 + mm_token_type_ids[(array_ids == self.image_token_id) & (~is_video_modality)] = 1 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) diff --git a/src/transformers/models/glm4v/processing_glm4v.py b/src/transformers/models/glm4v/processing_glm4v.py index c8fe12e39006..853a83fd9a23 100644 --- a/src/transformers/models/glm4v/processing_glm4v.py +++ b/src/transformers/models/glm4v/processing_glm4v.py @@ -36,7 +36,7 @@ class Glm4vProcessorKwargs(ProcessingKwargs, total=False): "text_kwargs": { "padding": False, "return_token_type_ids": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -58,6 +58,8 @@ def __init__(self, image_processor=None, tokenizer=None, video_processor=None, c else tokenizer.convert_tokens_to_ids(self.video_token) ) super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) + self.video_start_id = tokenizer.convert_tokens_to_ids("<|begin_of_video|>") + self.video_end_id = tokenizer.convert_tokens_to_ids("<|end_of_video|>") @auto_docstring def __call__( @@ -168,7 +170,16 @@ def __call__( if return_mm_token_type_ids: array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) - mm_token_type_ids[array_ids == self.image_token_id] = 1 + + # Replace 0 -> 2 only inside video segments because GLM4v + # uses the same special token to denote images and video + # Otherwise replace 0 -> 1 for image modality + starts = np.cumsum(array_ids == self.video_start_id, axis=1) + ends = np.cumsum(array_ids == self.video_end_id, axis=1) + is_video_modality = starts > ends + + mm_token_type_ids[(array_ids == self.image_token_id) & is_video_modality] = 2 + mm_token_type_ids[(array_ids == self.image_token_id) & (~is_video_modality)] = 1 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) @@ -237,6 +248,12 @@ def post_process_image_text_to_text( **kwargs, ) + @property + def model_input_names(self): + model_input_names = super().model_input_names + model_input_names.append("mm_token_type_ids") + return model_input_names + def replace_frame_token_id(self, timestamp_sec): return f"<|begin_of_image|>{self.image_token}<|end_of_image|>{int(timestamp_sec)}" diff --git a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py index 6edd4cedd4a4..4a012475a954 100644 --- a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py @@ -1121,51 +1121,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1180,93 +1226,71 @@ def get_rope_index( position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_start_token_id = self.config.video_start_token_id - video_end_token_id = self.config.video_end_token_id mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - video_group_index = 0 - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - - input_token_type = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False - if token == image_token_id and not video_check_flg: - input_token_type.append("image") - elif token == image_token_id and video_check_flg: - input_token_type.append("video") - else: - input_token_type.append("text") + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + video_group_index = 0 llm_pos_ids_list = [] - video_frame_num = 1 for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - if modality_type == "image": - t, h, w = image_grid_thw_list[image_index] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - image_index += 1 - video_frame_num = 1 - elif modality_type == "video": - _, h, w = video_grid_thw_list[video_index] - t = video_frame_num - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - video_group_index += 1 - if video_group_index >= video_grid_thw_list[video_index][0]: - video_index += 1 - video_group_index = 0 - video_frame_num += 1 + current_pos += text_len + # image == 1, video == 2 else: - text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - video_frame_num = 1 + # GLM4V_MOE splits video into segments per frame but there's only one `grid_thw` + # per whole video. We can't exhaus the iterator and have to re-use the grid + # while processing the same video! + if modality_type == 2: + if video_group_index == 0: + grid_thw = next(grid_iters[modality_type]) + video_group_index += 1 + video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index + else: + grid_thw = next(grid_iters[modality_type]) + + # Videos are processed per frame separately, each temporal grid is always `1` + temp_merge_size = grid_thw[0] + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1373,9 +1397,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1383,6 +1412,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1396,7 +1426,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1416,6 +1446,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Glm4vMoeModelOutputWithPast: @@ -1453,6 +1484,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1620,6 +1652,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1677,6 +1710,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + mm_token_type_ids=mm_token_type_ids, **kwargs, ) @@ -1771,8 +1805,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 4821bf829cd1..61bf605868a8 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -516,6 +516,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -531,6 +532,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + mm_token_type_ids=mm_token_type_ids, **kwargs, ) diff --git a/src/transformers/models/glm_image/modeling_glm_image.py b/src/transformers/models/glm_image/modeling_glm_image.py index 1a527d4ebb02..4639283fe8e5 100644 --- a/src/transformers/models/glm_image/modeling_glm_image.py +++ b/src/transformers/models/glm_image/modeling_glm_image.py @@ -974,6 +974,62 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, input_ids: torch.LongTensor | None = None, @@ -1050,9 +1106,6 @@ def get_rope_index( for img_idx, (start, end) in enumerate(zip(image_start_positions, image_end_positions)): if curr_grids is None or img_idx >= len(curr_grids): break - grid = curr_grids[img_idx] - # grid format is [temporal, height, width] - _, height, width = grid.tolist() # Text tokens before this image llm_pos_length = start - prev_image_end @@ -1063,14 +1116,10 @@ def get_rope_index( # For an image with height H and width W: # - position_width cycles [0, 1, ..., W-1] for each row, repeated H times # - position_height stays constant per row, [0]*W, [1]*W, ..., [H-1]*W - image_seq_length = height * width - position_width = torch.arange(current_pos, current_pos + width, device=device).repeat(height) - position_height = torch.arange(current_pos, current_pos + height, device=device).repeat_interleave( - width + vision_position_ids = self.get_vision_position_ids( + start_position=current_pos, grid_thw=curr_grids[img_idx], device=device ) - position_temporal = torch.full((image_seq_length,), current_pos, device=device, dtype=torch.long) - vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) - current_pos += max(height, width) + current_pos += max(curr_grids[img_idx][1], curr_grids[img_idx][2]) prev_image_end = end curr_position_ids.append(torch.cat([llm_position_ids, vision_position_ids], dim=-1)) diff --git a/src/transformers/models/glm_image/modular_glm_image.py b/src/transformers/models/glm_image/modular_glm_image.py index 5471b786db00..f9f207f1068a 100644 --- a/src/transformers/models/glm_image/modular_glm_image.py +++ b/src/transformers/models/glm_image/modular_glm_image.py @@ -737,9 +737,6 @@ def get_rope_index( for img_idx, (start, end) in enumerate(zip(image_start_positions, image_end_positions)): if curr_grids is None or img_idx >= len(curr_grids): break - grid = curr_grids[img_idx] - # grid format is [temporal, height, width] - _, height, width = grid.tolist() # Text tokens before this image llm_pos_length = start - prev_image_end @@ -750,14 +747,10 @@ def get_rope_index( # For an image with height H and width W: # - position_width cycles [0, 1, ..., W-1] for each row, repeated H times # - position_height stays constant per row, [0]*W, [1]*W, ..., [H-1]*W - image_seq_length = height * width - position_width = torch.arange(current_pos, current_pos + width, device=device).repeat(height) - position_height = torch.arange(current_pos, current_pos + height, device=device).repeat_interleave( - width + vision_position_ids = self.get_vision_position_ids( + start_position=current_pos, grid_thw=curr_grids[img_idx], device=device ) - position_temporal = torch.full((image_seq_length,), current_pos, device=device, dtype=torch.long) - vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) - current_pos += max(height, width) + current_pos += max(curr_grids[img_idx][1], curr_grids[img_idx][2]) prev_image_end = end curr_position_ids.append(torch.cat([llm_position_ids, vision_position_ids], dim=-1)) diff --git a/src/transformers/models/glm_ocr/modeling_glm_ocr.py b/src/transformers/models/glm_ocr/modeling_glm_ocr.py index d066f8428c10..8e37722918d4 100644 --- a/src/transformers/models/glm_ocr/modeling_glm_ocr.py +++ b/src/transformers/models/glm_ocr/modeling_glm_ocr.py @@ -867,51 +867,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -926,93 +972,71 @@ def get_rope_index( position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_start_token_id = self.config.video_start_token_id - video_end_token_id = self.config.video_end_token_id mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.ones( + position_ids = torch.zeros( 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - video_group_index = 0 - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - input_tokens = input_ids.tolist() - - input_token_type = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False - if token == image_token_id and not video_check_flg: - input_token_type.append("image") - elif token == image_token_id and video_check_flg: - input_token_type.append("video") - else: - input_token_type.append("text") + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + input_type_group = [] - for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]): + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): group = list(group) start_index = group[0][0] end_index = group[-1][0] + 1 input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + video_group_index = 0 llm_pos_ids_list = [] - video_frame_num = 1 for modality_type, start_idx, end_idx in input_type_group: - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - if modality_type == "image": - t, h, w = image_grid_thw_list[image_index] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - image_index += 1 - video_frame_num = 1 - elif modality_type == "video": - _, h, w = video_grid_thw_list[video_index] - t = video_frame_num - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - for t_idx in range(llm_grid_t): - t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx) - video_group_index += 1 - if video_group_index >= video_grid_thw_list[video_index][0]: - video_index += 1 - video_group_index = 0 - video_frame_num += 1 + current_pos += text_len + # image == 1, video == 2 else: - text_len = end_idx - start_idx - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - video_frame_num = 1 + # GLM_OCR splits video into segments per frame but there's only one `grid_thw` + # per whole video. We can't exhaus the iterator and have to re-use the grid + # while processing the same video! + if modality_type == 2: + if video_group_index == 0: + grid_thw = next(grid_iters[modality_type]) + video_group_index += 1 + video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index + else: + grid_thw = next(grid_iters[modality_type]) + + # Videos are processed per frame separately, each temporal grid is always `1` + temp_merge_size = grid_thw[0] + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1119,9 +1143,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1129,6 +1158,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1142,7 +1172,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1162,6 +1192,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | GlmOcrModelOutputWithPast: @@ -1199,6 +1230,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1311,6 +1343,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1363,6 +1396,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1448,8 +1482,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py index cd616538725c..4720dd398b3c 100644 --- a/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py @@ -23,6 +23,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -1060,44 +1061,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.embed_tokens = value + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [3, 4, 5, 6, 7] - text height position_ids: [3, 4, 5, 6, 7] - text width position_ids: [3, 4, 5, 6, 7] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1113,79 +1167,58 @@ def get_rope_index( mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) + mrope_position_deltas = [] position_ids = torch.zeros( - 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, ) - image_index, video_index = 0, 0 - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i].to(input_ids.device) == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = ( - image_grid_thw[image_index][0], - image_grid_thw[image_index][1], - image_grid_thw[image_index][2], + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - image_index += 1 - remain_images -= 1 - ed = ed_image + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = ( - video_grid_thw[video_index][0], - video_grid_thw[video_index][1], - video_grid_thw[video_index][2], + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device ) - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item(), - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(max(llm_positions.max() + 1 - len(total_input_ids[i]), 0)) - mrope_position_deltas = torch.tensor( - mrope_position_deltas, dtype=torch.long, device=input_ids.device - ).unsqueeze(1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @can_return_tuple @@ -1257,9 +1290,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1267,6 +1305,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1280,7 +1319,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1297,6 +1336,7 @@ def forward( use_cache: bool | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, rope_deltas: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs, @@ -1323,6 +1363,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1398,6 +1439,7 @@ def forward( pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1458,6 +1500,7 @@ def forward( use_cache=use_cache, pixel_values=pixel_values, rope_deltas=rope_deltas, + mm_token_type_ids=mm_token_type_ids, cache_position=cache_position, **kwargs, ) @@ -1539,8 +1582,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py index b6438baa7729..dc5bb2df1dcb 100644 --- a/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py @@ -471,6 +471,7 @@ class PaddleOCRVLProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, + "return_mm_token_type_ids": True, }, } @@ -493,6 +494,7 @@ class PaddleOCRVLProcessor(ProcessorMixin): def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id super().__init__(image_processor, tokenizer, chat_template=chat_template) def __call__( @@ -563,9 +565,17 @@ def __call__( index += 1 text[i] = text[i].replace("<|placeholder|>", self.image_token) - text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None) - return BatchFeature(data={**text_inputs, **image_inputs}) + if return_mm_token_type_ids: + array_ids = np.array(text_inputs["input_ids"]) + mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) + mm_token_type_ids[array_ids == self.image_token_id] = 1 + text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() + + return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors) class PaddleOCRVisionConfig(SiglipVisionConfig): @@ -1213,6 +1223,7 @@ def forward( use_cache: bool | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, rope_deltas: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs, @@ -1239,6 +1250,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1288,6 +1300,7 @@ def forward( pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1348,6 +1361,7 @@ def forward( use_cache=use_cache, pixel_values=pixel_values, rope_deltas=rope_deltas, + mm_token_type_ids=mm_token_type_ids, cache_position=cache_position, **kwargs, ) diff --git a/src/transformers/models/paddleocr_vl/processing_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/processing_paddleocr_vl.py index e33a23db4247..3f003364b847 100644 --- a/src/transformers/models/paddleocr_vl/processing_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/processing_paddleocr_vl.py @@ -23,6 +23,9 @@ # See the License for the specific language governing permissions and # limitations under the License. + +import numpy as np + from ...image_processing_utils import BatchFeature from ...image_utils import ImageInput from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack @@ -33,6 +36,7 @@ class PaddleOCRVLProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, + "return_mm_token_type_ids": True, }, } @@ -55,6 +59,7 @@ class PaddleOCRVLProcessor(ProcessorMixin): def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): self.image_token = tokenizer.image_token + self.image_token_id = tokenizer.image_token_id super().__init__(image_processor, tokenizer, chat_template=chat_template) def __call__( @@ -125,9 +130,17 @@ def __call__( index += 1 text[i] = text[i].replace("<|placeholder|>", self.image_token) - text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None) + + if return_mm_token_type_ids: + array_ids = np.array(text_inputs["input_ids"]) + mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) + mm_token_type_ids[array_ids == self.image_token_id] = 1 + text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() - return BatchFeature(data={**text_inputs, **image_inputs}) + return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors) __all__ = ["PaddleOCRVLProcessor"] diff --git a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py index 94e53202ff6d..6305902014ce 100644 --- a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -22,6 +22,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -1016,9 +1018,66 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, @@ -1026,42 +1085,31 @@ def get_rope_index( **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1079,14 +1127,9 @@ def get_rope_index( mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] + tokens_per_second = self.config.vision_config.tokens_per_second - total_input_ids = input_ids - if attention_mask is not None: - attention_mask = attention_mask == 1 + mrope_position_deltas = [] position_ids = torch.zeros( 3, input_ids.shape[0], @@ -1094,77 +1137,52 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - image_index, video_index = 0, 0 - for i, input_ids in enumerate(total_input_ids): + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + second_per_grid_ts = ( + iter(second_per_grid_ts) if second_per_grid_ts is not None else iter([1] * input_ids.shape[1]) + ) + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] if attention_mask is not None: - input_ids = input_ids[attention_mask[i]] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - second_per_grid_t = 0 - image_index += 1 - remain_images -= 1 - ed = ed_image + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - if second_per_grid_ts is not None: - second_per_grid_t = second_per_grid_ts[video_index] - else: - second_per_grid_t = 1.0 - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - range_tensor = torch.arange(llm_grid_t).view(-1, 1) - expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) - ## normalize type, send to device. - second_per_grid_t = torch.as_tensor( - second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device - ) - time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second - time_tensor_long = time_tensor.long() - t_index = time_tensor_long.flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + time_interval = tokens_per_second * int(next(second_per_grid_ts)) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, time_interval, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) if attention_mask is not None: - position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device) + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) else: - position_ids[..., i, :] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device) + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @can_return_tuple @@ -1261,9 +1279,14 @@ def compute_3d_position_ids( attention_mask: torch.Tensor | None, past_key_values: torch.Tensor | None, second_per_grid_ts: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1272,6 +1295,7 @@ def compute_3d_position_ids( video_grid_thw=video_grid_thw, attention_mask=attention_mask, second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1308,6 +1332,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], @@ -1357,6 +1382,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1484,6 +1510,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, logits_to_keep: int | torch.Tensor = 0, @@ -1552,6 +1579,7 @@ def forward( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1645,8 +1673,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py index c3ae8f885fed..093b4f2b783a 100644 --- a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py @@ -18,6 +18,8 @@ # limitations under the License. """PyTorch Qwen2.5-VL model.""" +import itertools + import numpy as np import torch import torch.nn as nn @@ -373,7 +375,8 @@ def __init__(self, config): def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, @@ -381,42 +384,31 @@ def get_rope_index( **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Temporal (Time): 3 patches, representing different segments of the video in time. - Height: 2 patches, dividing each frame vertically. - Width: 2 patches, dividing each frame horizontally. - We also have some important parameters: - fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. - 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. - temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. - 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. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [101, 102, 103, 104, 105] - text height position_ids: [101, 102, 103, 104, 105] - text width position_ids: [101, 102, 103, 104, 105] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -434,14 +426,9 @@ def get_rope_index( mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] + tokens_per_second = self.config.vision_config.tokens_per_second - total_input_ids = input_ids - if attention_mask is not None: - attention_mask = attention_mask == 1 + mrope_position_deltas = [] position_ids = torch.zeros( 3, input_ids.shape[0], @@ -449,77 +436,52 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - image_index, video_index = 0, 0 - for i, input_ids in enumerate(total_input_ids): + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + second_per_grid_ts = ( + iter(second_per_grid_ts) if second_per_grid_ts is not None else iter([1] * input_ids.shape[1]) + ) + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] if attention_mask is not None: - input_ids = input_ids[attention_mask[i]] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - second_per_grid_t = 0 - image_index += 1 - remain_images -= 1 - ed = ed_image + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - if second_per_grid_ts is not None: - second_per_grid_t = second_per_grid_ts[video_index] - else: - second_per_grid_t = 1.0 - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - range_tensor = torch.arange(llm_grid_t).view(-1, 1) - expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) - ## normalize type, send to device. - second_per_grid_t = torch.as_tensor( - second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device - ) - time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second - time_tensor_long = time_tensor.long() - t_index = time_tensor_long.flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + time_interval = tokens_per_second * int(next(second_per_grid_ts)) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, time_interval, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) if attention_mask is not None: - position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device) + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) else: - position_ids[..., i, :] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device) + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas def compute_3d_position_ids( @@ -531,9 +493,14 @@ def compute_3d_position_ids( attention_mask: torch.Tensor | None, past_key_values: torch.Tensor | None, second_per_grid_ts: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -542,6 +509,7 @@ def compute_3d_position_ids( video_grid_thw=video_grid_thw, attention_mask=attention_mask, second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -577,6 +545,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], @@ -626,6 +595,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -676,6 +646,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, second_per_grid_ts: torch.Tensor | None = None, logits_to_keep: int | torch.Tensor = 0, @@ -744,6 +715,7 @@ def forward( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -824,7 +796,7 @@ class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -839,7 +811,7 @@ def model_input_names(self): names_from_processor = list( dict.fromkeys(tokenizer_input_names + image_processor_input_names + video_processor_input_names) ) - return names_from_processor + ["second_per_grid_ts"] + return names_from_processor + ["second_per_grid_ts", "mm_token_type_ids"] def __call__( self, @@ -928,6 +900,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) diff --git a/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py index 7b94fe8bd196..6082653751e1 100644 --- a/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/processing_qwen2_5_vl.py @@ -37,7 +37,7 @@ class Qwen2_5_VLProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -148,6 +148,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) @@ -225,7 +226,7 @@ def model_input_names(self): names_from_processor = list( dict.fromkeys(tokenizer_input_names + image_processor_input_names + video_processor_input_names) ) - return names_from_processor + ["second_per_grid_ts"] + return names_from_processor + ["second_per_grid_ts", "mm_token_type_ids"] __all__ = ["Qwen2_5_VLProcessor"] diff --git a/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py index 16c8962d0058..3a403942fc18 100644 --- a/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -18,6 +18,7 @@ # limitations under the License. """PyTorch Qwen2-VL model.""" +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -990,44 +991,97 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ - Calculate the 3D rope index based on image and video's temporal, height and width in LLM. - - Explanation: - Each embedding sequence contains vision embedding and text embedding or just contains text embedding. - - For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. - Examples: - input_ids: [T T T T T], here T is for text. - temporal position_ids: [0, 1, 2, 3, 4] - height position_ids: [0, 1, 2, 3, 4] - width position_ids: [0, 1, 2, 3, 4] - - For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part - and 1D rotary position embedding for text part. - Examples: - Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches. - input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. - vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2] - vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] - vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] - text temporal position_ids: [3, 4, 5, 6, 7] - text height position_ids: [3, 4, 5, 6, 7] - text width position_ids: [3, 4, 5, 6, 7] - Here we calculate the text start position_ids as the max vision position_ids plus 1. + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` + + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` + + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): @@ -1043,79 +1097,58 @@ def get_rope_index( mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) + mrope_position_deltas = [] position_ids = torch.zeros( - 3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, ) - image_index, video_index = 0, 0 - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i].to(input_ids.device) == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = ( - image_grid_thw[image_index][0], - image_grid_thw[image_index][1], - image_grid_thw[image_index][2], + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos ) - image_index += 1 - remain_images -= 1 - ed = ed_image + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = ( - video_grid_thw[video_index][0], - video_grid_thw[video_index][1], - video_grid_thw[video_index][2], + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device ) - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t.item(), - h.item() // spatial_merge_size, - w.item() // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(max(llm_positions.max() + 1 - len(total_input_ids[i]), 0)) - mrope_position_deltas = torch.tensor( - mrope_position_deltas, dtype=torch.long, device=input_ids.device - ).unsqueeze(1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @can_return_tuple @@ -1211,9 +1244,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1221,6 +1259,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1234,7 +1273,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1257,6 +1296,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen2VLModelOutputWithPast: @@ -1302,6 +1342,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1398,6 +1439,7 @@ def forward( image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, rope_deltas: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1462,6 +1504,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1552,8 +1595,10 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen2_vl/processing_qwen2_vl.py b/src/transformers/models/qwen2_vl/processing_qwen2_vl.py index 3cd29ad0e610..bcb9ac383154 100644 --- a/src/transformers/models/qwen2_vl/processing_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/processing_qwen2_vl.py @@ -37,7 +37,7 @@ class Qwen2VLProcessorKwargs(ProcessingKwargs, total=False): _defaults = { "text_kwargs": { "padding": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, } @@ -129,6 +129,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) @@ -198,5 +199,11 @@ def post_process_image_text_to_text( **kwargs, ) + @property + def model_input_names(self): + model_input_names = super().model_input_names + model_input_names.append("mm_token_type_ids") + return model_input_names + __all__ = ["Qwen2VLProcessor"] diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 327e4003479f..2a49fcb82d67 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -1423,32 +1424,114 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: - """Different from the original implementation, Qwen3_5 use timestamps rather than absolute time position ids.""" + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` - # Since we use timestamps to separate videos, like , the video_grid_thw should also be split - if video_grid_thw is not None: - video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - video_grid_thw[:, 0] = 1 + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) position_ids = torch.zeros( 3, input_ids.shape[0], @@ -1456,59 +1539,48 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1603,9 +1675,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1613,6 +1690,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1626,7 +1704,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1645,6 +1723,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3_5ModelOutputWithPast: @@ -1690,6 +1769,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1885,6 +1965,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1947,6 +2028,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + mm_token_type_ids=mm_token_type_ids, **kwargs, ) @@ -2019,16 +2101,18 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): if (cache := model_kwargs.get("past_key_values")) is not None: past_length = cache.get_seq_length() if past_length != 0 and self.model.rope_deltas is not None: - text_positions += self.model.rope_deltas - return text_positions + position_ids = text_positions[None, ...] + self.model.rope_deltas + return position_ids # Otherwise compute 3d position ids for vision tokens and concat with text position ids if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 8e2be547d0d2..cae6a0a7d383 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -731,6 +731,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3_5ModelOutputWithPast: @@ -776,6 +777,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 4976abb5447b..87f52dc8b651 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -1548,32 +1549,114 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: - """Different from the original implementation, Qwen3_5Moe use timestamps rather than absolute time position ids.""" + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` - # Since we use timestamps to separate videos, like , the video_grid_thw should also be split - if video_grid_thw is not None: - video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - video_grid_thw[:, 0] = 1 + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) position_ids = torch.zeros( 3, input_ids.shape[0], @@ -1581,59 +1664,48 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1728,9 +1800,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1738,6 +1815,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1751,7 +1829,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1770,6 +1848,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3_5MoeModelOutputWithPast: @@ -1815,6 +1894,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -2087,6 +2167,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -2148,6 +2229,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -2240,16 +2322,18 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): if (cache := model_kwargs.get("past_key_values")) is not None: past_length = cache.get_seq_length() if past_length != 0 and self.model.rope_deltas is not None: - text_positions += self.model.rope_deltas - return text_positions + position_ids = text_positions[None, ...] + self.model.rope_deltas + return position_ids # Otherwise compute 3d position ids for vision tokens and concat with text position ids if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index 2a3e2e52fe5b..8e5ab18be6a2 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -977,32 +978,114 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: - """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids.""" + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` - # Since we use timestamps to separate videos, like , the video_grid_thw should also be split - if video_grid_thw is not None: - video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - video_grid_thw[:, 0] = 1 + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) position_ids = torch.zeros( 3, input_ids.shape[0], @@ -1010,59 +1093,48 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1157,9 +1229,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1167,6 +1244,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1180,7 +1258,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1199,6 +1277,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3VLModelOutputWithPast: @@ -1273,6 +1352,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1387,6 +1467,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1449,6 +1530,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + mm_token_type_ids=mm_token_type_ids, **kwargs, ) @@ -1521,16 +1603,18 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): if (cache := model_kwargs.get("past_key_values")) is not None: past_length = cache.get_seq_length() if past_length != 0 and self.model.rope_deltas is not None: - text_positions += self.model.rope_deltas - return text_positions + position_ids = text_positions[None, ...] + self.model.rope_deltas + return position_ids # Otherwise compute 3d position ids for vision tokens and concat with text position ids if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py index 9eed1b6ff890..e791940ea965 100644 --- a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py @@ -852,95 +852,6 @@ def __init__(self, config): self.visual = Qwen3VLVisionModel._from_config(config.vision_config) self.language_model = Qwen3VLTextModel._from_config(config.text_config) - def get_rope_index( - self, - input_ids: torch.LongTensor | None = None, - image_grid_thw: torch.LongTensor | None = None, - video_grid_thw: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Different from the original implementation, Qwen3VL use timestamps rather than absolute time position ids.""" - - # Since we use timestamps to separate videos, like , the video_grid_thw should also be split - if video_grid_thw is not None: - video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - video_grid_thw[:, 0] = 1 - - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None - - spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id - mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) - position_ids = torch.zeros( - 3, - input_ids.shape[0], - input_ids.shape[1], - dtype=input_ids.dtype, - device=input_ids.device, - ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image - else: - t, h, w = video_grid_thw_list[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) - return position_ids, mrope_position_deltas - @can_return_tuple @auto_docstring def get_image_features( @@ -996,6 +907,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3VLModelOutputWithPast: @@ -1070,6 +982,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1119,6 +1032,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1181,6 +1095,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + mm_token_type_ids=mm_token_type_ids, **kwargs, ) @@ -1243,42 +1158,6 @@ def prepare_inputs_for_generation( return model_inputs - def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): - # Overwritten -- requires 3D position ids - - text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs) - - # Early exit in case we are continuing generation from past kv - past_length = 0 - if (cache := model_kwargs.get("past_key_values")) is not None: - past_length = cache.get_seq_length() - if past_length != 0 and self.model.rope_deltas is not None: - text_positions += self.model.rope_deltas - return text_positions - - # Otherwise compute 3d position ids for vision tokens and concat with text position ids - if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: - inputs_tensor = model_kwargs["input_ids"] - - is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None - ): - model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} - vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) - self.model.rope_deltas = rope_deltas - else: - vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1) - self.model.rope_deltas = torch.zeros( - inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device - ) - - # Concatenate "text + vision" positions into [4, bs, seq-len] - text_positions = text_positions[None, ...] - position_ids = torch.cat([text_positions, vision_positions], dim=0) - - return position_ids - def _expand_inputs_for_generation( self, expand_size: int = 1, @@ -1382,7 +1261,7 @@ class Qwen3VLProcessorKwargs(ProcessingKwargs, total=False): "text_kwargs": { "padding": False, "return_token_type_ids": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -1515,6 +1394,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) diff --git a/src/transformers/models/qwen3_vl/processing_qwen3_vl.py b/src/transformers/models/qwen3_vl/processing_qwen3_vl.py index 29fa0231daef..e25ecbda4b7f 100644 --- a/src/transformers/models/qwen3_vl/processing_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/processing_qwen3_vl.py @@ -36,7 +36,7 @@ class Qwen3VLProcessorKwargs(ProcessingKwargs, total=False): "text_kwargs": { "padding": False, "return_token_type_ids": False, - "return_mm_token_type_ids": False, + "return_mm_token_type_ids": True, }, "videos_kwargs": {"return_metadata": True}, } @@ -183,6 +183,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) @@ -252,6 +253,12 @@ def post_process_image_text_to_text( **kwargs, ) + @property + def model_input_names(self): + model_input_names = super().model_input_names + model_input_names.append("mm_token_type_ids") + return model_input_names + def _calculate_timestamps(self, indices: list[int] | np.ndarray, video_fps: float, merge_size: int = 2): if not isinstance(indices, list): indices = indices.tolist() diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 354bd52997fe..b3d8e371f84a 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools from collections.abc import Callable from dataclasses import dataclass from typing import Any, Optional @@ -1106,32 +1107,114 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + def get_vision_position_ids( + self, + start_position: int, + grid_thw: list[int, int, int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ): + """ + Compute 3D positional indices for vision tokens derived from a single image or video input. + + The positions are generated from the input grid defined by temporal (T), height (H), and + width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the + merge sizes used in the vision backbone. The resulting positions are offset by `start_position`. + + Args: + start_position (`int`): + Offset added to all computed positional indices. + grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`): + The (T, H, W) grid representing the feature layout of the current image or video after patch embedding. + temp_merge_size (`int`, *optional*): + Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided + by this value. Defaults to 1. + spatial_merge_size (`int`, *optional*): + Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided + by this value. Defaults to 1. + time_interval (`int`, *optional*): + Spacing factor applied between consecutive temporal position indices.Defaults to 1. + device (`str` or `torch.device`, *optional*): + Device on which the resulting tensor is allocated. If `None`, uses the current default device. + + Returns: + torch.LongTensor of shape (3, sequence_length): + Positional indices for temporal, height, and width dimensions, + flattened into sequence form and offset by `start_position`. + """ + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat( + llm_grid_h * llm_grid_t + ) + position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave( + llm_grid_w * llm_grid_t + ) + position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long) + position_temporal = position_temporal * time_interval + vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0) + + return vision_position_ids + def get_rope_index( self, - input_ids: torch.LongTensor | None = None, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.IntTensor, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: - """Different from the original implementation, Qwen3VLMoe use timestamps rather than absolute time position ids.""" + """ + Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text` + sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred + position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width) + while text tokens use standard 1D RoPE. + + Example: + Temporal patches: 3; Height patches: 2; Width patches: 2 + Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total. + + Temporal position IDs are spaced by: + `interval = tokens_per_second * temporal_patch_size / fps` - # Since we use timestamps to separate videos, like , the video_grid_thw should also be split - if video_grid_thw is not None: - video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) - video_grid_thw[:, 0] = 1 + If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch: + `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]` - image_grid_thw_list = image_grid_thw.tolist() if image_grid_thw is not None else None - video_grid_thw_list = video_grid_thw.tolist() if video_grid_thw is not None else None + Height IDs repeat per row: `[0, 0, 1, 1, ...]` + Width IDs alternate per column: `[0, 1, 0, 1, ...]` + Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1` + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`): + Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2). + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ spatial_merge_size = self.config.vision_config.spatial_merge_size - image_token_id = self.config.image_token_id - video_token_id = self.config.video_token_id - vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] - total_input_ids = input_ids - if attention_mask is None: - attention_mask = torch.ones_like(total_input_ids) position_ids = torch.zeros( 3, input_ids.shape[0], @@ -1139,59 +1222,48 @@ def get_rope_index( dtype=input_ids.dtype, device=input_ids.device, ) - image_index, video_index = 0, 0 - attention_mask = attention_mask.to(total_input_ids.device) - for i, input_ids in enumerate(total_input_ids): - input_ids = input_ids[attention_mask[i] == 1] - image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) - vision_tokens = input_ids[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() - input_tokens = input_ids.tolist() - llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - for _ in range(image_nums + video_nums): - if image_token_id in input_tokens and remain_images > 0: - ed_image = input_tokens.index(image_token_id, st) - else: - ed_image = len(input_tokens) + 1 - if video_token_id in input_tokens and remain_videos > 0: - ed_video = input_tokens.index(video_token_id, st) - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw_list[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + # text == 0 + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + # image == 1, video == 2 else: - t, h, w = video_grid_thw_list[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) - # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = self.get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) return position_ids, mrope_position_deltas @@ -1286,9 +1358,14 @@ def compute_3d_position_ids( video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, past_key_values: torch.Tensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, ) -> torch.Tensor | None: past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length() - can_compute_mrope = input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None) + can_compute_mrope = ( + input_ids is not None + and mm_token_type_ids is not None + and (image_grid_thw is not None or video_grid_thw is not None) + ) if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0): position_ids, rope_deltas = self.get_rope_index( @@ -1296,6 +1373,7 @@ def compute_3d_position_ids( image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask, + mm_token_type_ids=mm_token_type_ids, ) self.rope_deltas = rope_deltas # Use pre-calculated rope-deltas to infer correct 3D position ids @@ -1309,7 +1387,7 @@ def compute_3d_position_ids( position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length) position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device) delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0) - position_ids = position_ids + delta.to(device=position_ids.device) + position_ids = position_ids + delta.to(device=inputs_embeds.device) else: # Can't build correct 3D positions. Let the model infer it from `cache_position` position_ids = None @@ -1328,6 +1406,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen3VLMoeModelOutputWithPast: @@ -1402,6 +1481,7 @@ def forward( inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, + mm_token_type_ids=mm_token_type_ids, ) outputs = self.language_model( @@ -1569,6 +1649,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -1631,6 +1712,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, @@ -1723,16 +1805,18 @@ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): if (cache := model_kwargs.get("past_key_values")) is not None: past_length = cache.get_seq_length() if past_length != 0 and self.model.rope_deltas is not None: - text_positions += self.model.rope_deltas - return text_positions + position_ids = text_positions[None, ...] + self.model.rope_deltas + return position_ids # Otherwise compute 3d position ids for vision tokens and concat with text position ids if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: inputs_tensor = model_kwargs["input_ids"] is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] - if is_input_ids and ( - model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None + if ( + is_input_ids + and model_kwargs.get("mm_token_type_ids") is not None + and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None) ): model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"} vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs) diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index 203b5fad9c19..6033e30bbfb6 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -458,6 +458,7 @@ def forward( pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.IntTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -520,6 +521,7 @@ def forward( pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + mm_token_type_ids=mm_token_type_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, diff --git a/src/transformers/models/video_llama_3/modular_video_llama_3.py b/src/transformers/models/video_llama_3/modular_video_llama_3.py index 31083204afee..9ea9b1fcb370 100644 --- a/src/transformers/models/video_llama_3/modular_video_llama_3.py +++ b/src/transformers/models/video_llama_3/modular_video_llama_3.py @@ -596,6 +596,9 @@ def __init__(self, config: VideoLlama3Config): def get_rope_index(self): raise AttributeError("Not needed for VideoLLaMA3") + def get_vision_position_ids(self): + raise AttributeError("Not needed for VideoLLaMA3") + def compute_3d_position_ids(self): raise AttributeError("Not needed for VideoLLaMA3") @@ -1202,10 +1205,14 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) + def model_input_names(self): + raise AttributeError("VideoLlama doesn't need to override it") + class VideoLlama3ImageProcessorKwargs(Qwen2VLImageProcessorKwargs): pass diff --git a/src/transformers/models/video_llama_3/processing_video_llama_3.py b/src/transformers/models/video_llama_3/processing_video_llama_3.py index 14ac00bf9c93..0bfbb76757c3 100644 --- a/src/transformers/models/video_llama_3/processing_video_llama_3.py +++ b/src/transformers/models/video_llama_3/processing_video_llama_3.py @@ -156,6 +156,7 @@ def __call__( array_ids = np.array(text_inputs["input_ids"]) mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) mm_token_type_ids[array_ids == self.image_token_id] = 1 + mm_token_type_ids[array_ids == self.video_token_id] = 2 text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) diff --git a/src/transformers/utils/auto_docstring.py b/src/transformers/utils/auto_docstring.py index 02613cd53922..ca9841158219 100644 --- a/src/transformers/utils/auto_docstring.py +++ b/src/transformers/utils/auto_docstring.py @@ -613,6 +613,15 @@ class ModelArgs: "shape": "of shape `(batch_size, sequence_length)`", } + mm_token_type_ids = { + "description": """ + Indices of input sequence tokens matching each modality. For example text (0), image (1), video (2). + Multimodal token type ids can be obtained using [`AutoProcessor`]. See [`ProcessorMixin.__call__`] for details. + + """, + "shape": "of shape `(batch_size, sequence_length)`", + } + position_ids = { "description": """ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. diff --git a/tests/models/glm46v/test_modeling_glm46v.py b/tests/models/glm46v/test_modeling_glm46v.py index 80522bf08e77..eb5489b0e1a2 100644 --- a/tests/models/glm46v/test_modeling_glm46v.py +++ b/tests/models/glm46v/test_modeling_glm46v.py @@ -159,6 +159,9 @@ def prepare_config_and_inputs_for_common(self): patch_size = config.vision_config.patch_size patches_per_side = self.image_size // patch_size + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, 1 : 1 + self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor( @@ -166,6 +169,7 @@ def prepare_config_and_inputs_for_common(self): ), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict diff --git a/tests/models/glm46v/test_processor_glm46v.py b/tests/models/glm46v/test_processor_glm46v.py index 344e2e293727..8e7b56df6fa8 100644 --- a/tests/models/glm46v/test_processor_glm46v.py +++ b/tests/models/glm46v/test_processor_glm46v.py @@ -185,7 +185,7 @@ def test_apply_chat_template_video_frame_sampling(self): self.assertListEqual(expected_output, formatted_prompt_tokenized) out_dict = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True) - self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask"]) + self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask", "mm_token_type_ids"]) # Add video URL for return dict and load with `num_frames` arg messages[0][0]["content"][0] = { diff --git a/tests/models/glm4v/test_modeling_glm4v.py b/tests/models/glm4v/test_modeling_glm4v.py index d1d247cace2e..c96547d883a5 100644 --- a/tests/models/glm4v/test_modeling_glm4v.py +++ b/tests/models/glm4v/test_modeling_glm4v.py @@ -159,6 +159,9 @@ def prepare_config_and_inputs_for_common(self): patch_size = config.vision_config.patch_size patches_per_side = self.image_size // patch_size + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, 1 : 1 + self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor( @@ -166,6 +169,7 @@ def prepare_config_and_inputs_for_common(self): ), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict diff --git a/tests/models/glm4v/test_processor_glm4v.py b/tests/models/glm4v/test_processor_glm4v.py index 5acf39e6e731..cb101521ea24 100644 --- a/tests/models/glm4v/test_processor_glm4v.py +++ b/tests/models/glm4v/test_processor_glm4v.py @@ -185,7 +185,7 @@ def test_apply_chat_template_video_frame_sampling(self): self.assertListEqual(expected_output, formatted_prompt_tokenized) out_dict = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True) - self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask"]) + self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask", "mm_token_type_ids"]) # Add video URL for return dict and load with `num_frames` arg messages[0][0]["content"][0] = { diff --git a/tests/models/glm4v_moe/test_modeling_glm4v_moe.py b/tests/models/glm4v_moe/test_modeling_glm4v_moe.py index 8572dc1e229f..917a13bcefaf 100644 --- a/tests/models/glm4v_moe/test_modeling_glm4v_moe.py +++ b/tests/models/glm4v_moe/test_modeling_glm4v_moe.py @@ -168,6 +168,9 @@ def prepare_config_and_inputs_for_common(self): patch_size = config.vision_config.patch_size patches_per_side = self.image_size // patch_size + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, 1 : 1 + self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor( @@ -175,6 +178,7 @@ def prepare_config_and_inputs_for_common(self): ), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict diff --git a/tests/models/glm_ocr/test_modeling_glm_ocr.py b/tests/models/glm_ocr/test_modeling_glm_ocr.py index 7f3dfdaed437..0d0a809ec983 100644 --- a/tests/models/glm_ocr/test_modeling_glm_ocr.py +++ b/tests/models/glm_ocr/test_modeling_glm_ocr.py @@ -170,6 +170,9 @@ def prepare_config_and_inputs_for_common(self): patch_size = config.vision_config.patch_size patches_per_side = self.image_size // patch_size + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, 1 : 1 + self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor( @@ -177,6 +180,7 @@ def prepare_config_and_inputs_for_common(self): ), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict diff --git a/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py b/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py index 57288dd01f2e..a59415e378a4 100644 --- a/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py +++ b/tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py @@ -124,6 +124,7 @@ def __init__( "spatial_patch_size": 14, "spatial_merge_size": 1, "temporal_patch_size": 2, + "tokens_per_second": 1, } self.vision_config = vision_config self.text_config = { @@ -178,11 +179,16 @@ def prepare_config_and_inputs_for_common(self): input_ids[input_ids == self.vision_start_token_id] = self.pad_token_id input_ids[:, self.num_image_tokens] = self.image_token_id input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id + + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict @@ -230,6 +236,7 @@ def test_mismatching_num_image_tokens(self): with self.assertRaisesRegex(ValueError, "Image features and image tokens do not match"): _ = model(**curr_input_dict) + model.base_model.rope_deltas = None # simulate multi-image case by concatenating inputs where each has exactly one image/image-token input_ids = curr_input_dict["input_ids"][:1] pixel_values = curr_input_dict["pixel_values"][:one_img_length] @@ -244,6 +251,7 @@ def test_mismatching_num_image_tokens(self): image_grid_thw=image_grid_thw, ) + model.base_model.rope_deltas = None # two images and two image tokens don't raise an error pixel_values = torch.cat([pixel_values, pixel_values], dim=0) image_grid_thw = torch.cat([image_grid_thw, image_grid_thw], dim=0) @@ -353,6 +361,7 @@ def attention_mask_padding_matches_padding_free_with_position_ids( input_ids=inputs_dict["input_ids"], image_grid_thw=inputs_dict["image_grid_thw"], attention_mask=inputs_dict["attention_mask"], + mm_token_type_ids=inputs_dict["mm_token_type_ids"], ) # [3, bs, padded-seq-len] vision_padfree_positions = vision_position_ids[:, dummy_attention_mask.bool()].view( 3, -1 diff --git a/tests/models/qwen2_5_vl/test_processing_qwen2_5_vl.py b/tests/models/qwen2_5_vl/test_processing_qwen2_5_vl.py index ee182c15005c..c0f4b7240fb8 100644 --- a/tests/models/qwen2_5_vl/test_processing_qwen2_5_vl.py +++ b/tests/models/qwen2_5_vl/test_processing_qwen2_5_vl.py @@ -189,7 +189,7 @@ def test_apply_chat_template_video_frame_sampling(self): self.assertListEqual(expected_output, formatted_prompt_tokenized) out_dict = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True) - self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask"]) + self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask", "mm_token_type_ids"]) # Add video URL for return dict and load with `num_frames` arg messages[0][0]["content"][0] = { diff --git a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py b/tests/models/qwen2_vl/test_modeling_qwen2_vl.py index 02d82ca6217d..dbc476596a36 100644 --- a/tests/models/qwen2_vl/test_modeling_qwen2_vl.py +++ b/tests/models/qwen2_vl/test_modeling_qwen2_vl.py @@ -156,11 +156,15 @@ def prepare_config_and_inputs_for_common(self): input_ids[:, self.num_image_tokens] = self.image_token_id input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict @@ -213,6 +217,7 @@ def test_mismatching_num_image_tokens(self): with self.assertRaisesRegex(ValueError, "Image features and image tokens do not match"): _ = model(**curr_input_dict) + model.base_model.rope_deltas = None # simulate multi-image case by concatenating inputs where each has exactly one image/image-token input_ids = curr_input_dict["input_ids"][:1] pixel_values = curr_input_dict["pixel_values"][:one_img_length] @@ -223,6 +228,7 @@ def test_mismatching_num_image_tokens(self): with self.assertRaisesRegex(ValueError, "Image features and image tokens do not match"): _ = model(input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw) + model.base_model.rope_deltas = None # two images and two image tokens don't raise an error pixel_values = torch.cat([pixel_values, pixel_values], dim=0) image_grid_thw = torch.cat([image_grid_thw, image_grid_thw], dim=0) @@ -306,6 +312,7 @@ def attention_mask_padding_matches_padding_free_with_position_ids( input_ids=inputs_dict["input_ids"], image_grid_thw=inputs_dict["image_grid_thw"], attention_mask=inputs_dict["attention_mask"], + mm_token_type_ids=inputs_dict["mm_token_type_ids"], ) # [3, bs, padded-seq-len] vision_padfree_positions = vision_position_ids[:, dummy_attention_mask.bool()].view( 3, -1 diff --git a/tests/models/qwen2_vl/test_processing_qwen2_vl.py b/tests/models/qwen2_vl/test_processing_qwen2_vl.py index b2a25eafb3a5..db5236573c85 100644 --- a/tests/models/qwen2_vl/test_processing_qwen2_vl.py +++ b/tests/models/qwen2_vl/test_processing_qwen2_vl.py @@ -191,7 +191,7 @@ def test_apply_chat_template_video_frame_sampling(self): self.assertListEqual(expected_output, formatted_prompt_tokenized) out_dict = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True) - self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask"]) + self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask", "mm_token_type_ids"]) # Add video URL for return dict and load with `num_frames` arg messages[0][0]["content"][0] = { diff --git a/tests/models/qwen3_5/test_modeling_qwen3_5.py b/tests/models/qwen3_5/test_modeling_qwen3_5.py index da025439acb4..2e844db29c94 100644 --- a/tests/models/qwen3_5/test_modeling_qwen3_5.py +++ b/tests/models/qwen3_5/test_modeling_qwen3_5.py @@ -277,11 +277,16 @@ def prepare_config_and_inputs_for_common(self): input_ids[input_ids == self.vision_start_token_id] = self.pad_token_id input_ids[:, self.num_image_tokens] = self.image_token_id input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id + + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict diff --git a/tests/models/qwen3_vl/test_modeling_qwen3_vl.py b/tests/models/qwen3_vl/test_modeling_qwen3_vl.py index c0c8579d1b2e..91692b386b87 100644 --- a/tests/models/qwen3_vl/test_modeling_qwen3_vl.py +++ b/tests/models/qwen3_vl/test_modeling_qwen3_vl.py @@ -164,11 +164,16 @@ def prepare_config_and_inputs_for_common(self): input_ids[input_ids == self.vision_start_token_id] = self.pad_token_id input_ids[:, self.num_image_tokens] = self.image_token_id input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id + + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict @@ -216,6 +221,7 @@ def test_mismatching_num_image_tokens(self): with self.assertRaises(ValueError): _ = model(**curr_input_dict) + model.base_model.rope_deltas = None # simulate multi-image case by concatenating inputs where each has exactly one image/image-token input_ids = curr_input_dict["input_ids"][:1] pixel_values = curr_input_dict["pixel_values"][:one_img_length] @@ -230,6 +236,7 @@ def test_mismatching_num_image_tokens(self): image_grid_thw=image_grid_thw, ) + model.base_model.rope_deltas = None # two images and two image tokens don't raise an error pixel_values = torch.cat([pixel_values, pixel_values], dim=0) image_grid_thw = torch.cat([image_grid_thw, image_grid_thw], dim=0) diff --git a/tests/models/qwen3_vl/test_processing_qwen3_vl.py b/tests/models/qwen3_vl/test_processing_qwen3_vl.py index 763d445c5306..487f3618a46f 100644 --- a/tests/models/qwen3_vl/test_processing_qwen3_vl.py +++ b/tests/models/qwen3_vl/test_processing_qwen3_vl.py @@ -210,7 +210,7 @@ def test_apply_chat_template_video_frame_sampling(self): self.assertListEqual(expected_output, formatted_prompt_tokenized) out_dict = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True) - self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask"]) + self.assertListEqual(list(out_dict.keys()), ["input_ids", "attention_mask", "mm_token_type_ids"]) # for fast test, set the longest edge to 8192 processor.video_processor.size["longest_edge"] = 8192 diff --git a/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py b/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py index a5c8d3e4d953..eff30c520c8f 100644 --- a/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py +++ b/tests/models/qwen3_vl_moe/test_modeling_qwen3_vl_moe.py @@ -168,11 +168,16 @@ def prepare_config_and_inputs_for_common(self): input_ids[input_ids == self.vision_start_token_id] = self.pad_token_id input_ids[:, self.num_image_tokens] = self.image_token_id input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id + + mm_token_type_ids = torch.zeros_like(input_ids) + mm_token_type_ids[:, self.num_image_tokens] = 1 + inputs_dict = { "pixel_values": pixel_values, "image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device), "input_ids": input_ids, "attention_mask": attention_mask, + "mm_token_type_ids": mm_token_type_ids, } return config, inputs_dict @@ -220,6 +225,7 @@ def test_mismatching_num_image_tokens(self): with self.assertRaisesRegex(ValueError, "Image features and image tokens do not match"): _ = model(**curr_input_dict) + model.base_model.rope_deltas = None # simulate multi-image case by concatenating inputs where each has exactly one image/image-token input_ids = curr_input_dict["input_ids"][:1] pixel_values = curr_input_dict["pixel_values"][:one_img_length] @@ -234,6 +240,7 @@ def test_mismatching_num_image_tokens(self): image_grid_thw=image_grid_thw, ) + model.base_model.rope_deltas = None # two images and two image tokens don't raise an error pixel_values = torch.cat([pixel_values, pixel_values], dim=0) image_grid_thw = torch.cat([image_grid_thw, image_grid_thw], dim=0)