# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/fgclip2/modular_fgclip2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_fgclip2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 # Copyright 2025 The HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math import warnings from dataclasses import dataclass from typing import Any, Callable, Optional, Union, List import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.init import _calculate_fan_in_and_fan_out from torchvision.ops import roi_align from transformers.activations import ACT2FN from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.processing_utils import Unpack from transformers.utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs from transformers.utils.generic import check_model_inputs from .configuration_fgclip2 import Fgclip2Config, Fgclip2TextConfig, Fgclip2VisionConfig @dataclass @auto_docstring( custom_intro=""" Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. """ ) class Fgclip2VisionOutput(ModelOutput): r""" image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): The image embeddings obtained by applying the projection layer to the pooler_output. """ image_embeds: Optional[torch.FloatTensor] = None last_hidden_state: Optional[torch.FloatTensor] = None hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None attentions: Optional[tuple[torch.FloatTensor, ...]] = None @dataclass @auto_docstring( custom_intro=""" Base class for text model's outputs that also contains a pooling of the last hidden states. """ ) class Fgclip2TextOutput(ModelOutput): r""" text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): The text embeddings obtained by applying the projection layer to the pooler_output. """ text_embeds: Optional[torch.FloatTensor] = None last_hidden_state: Optional[torch.FloatTensor] = None hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None attentions: Optional[tuple[torch.FloatTensor, ...]] = None @dataclass @auto_docstring class Fgclip2Output(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Contrastive loss for image-text similarity. logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text similarity scores. logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image similarity scores. text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`Fgclip2TextModel`]. image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`Fgclip2VisionModel`]. text_model_output (`BaseModelOutputWithPooling`): The output of the [`Fgclip2TextModel`]. vision_model_output (`BaseModelOutputWithPooling`): The output of the [`Fgclip2VisionModel`]. """ loss: Optional[torch.FloatTensor] = None logits_per_image: Optional[torch.FloatTensor] = None logits_per_text: Optional[torch.FloatTensor] = None text_embeds: Optional[torch.FloatTensor] = None image_embeds: Optional[torch.FloatTensor] = None text_model_output: BaseModelOutputWithPooling = None vision_model_output: BaseModelOutputWithPooling = None def to_tuple(self) -> tuple[Any]: return tuple( self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple() for k in self.keys() ) class Fgclip2VisionEmbeddings(nn.Module): def __init__(self, config: Fgclip2VisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.patch_size = config.patch_size self.patch_embedding = nn.Linear( in_features=config.num_channels * self.patch_size * self.patch_size, out_features=self.embed_dim, ) self.num_patches = config.num_patches self.position_embedding_size = int(self.num_patches**0.5) self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim) @staticmethod def resize_positional_embeddings( positional_embeddings: torch.Tensor, spatial_shapes: torch.LongTensor, max_length: int, ) -> torch.Tensor: """ Resize positional embeddings to image-specific size and pad to a fixed size. Args: positional_embeddings (`torch.Tensor`): Position embeddings of shape (height, width, embed_dim) spatial_shapes (`torch.LongTensor`): Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to max_length (`int`): Maximum length of the positional embeddings to pad resized positional embeddings to Returns: `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim) """ batch_size = spatial_shapes.shape[0] embed_dim = positional_embeddings.shape[-1] source_dtype = positional_embeddings.dtype resulted_positional_embeddings = torch.empty( (batch_size, max_length, embed_dim), device=positional_embeddings.device, dtype=source_dtype, ) # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0) # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU if positional_embeddings.device.type == "cpu": positional_embeddings = positional_embeddings.to(torch.float32) for i in range(batch_size): # (1, dim, height, width) -> (1, dim, target_height, target_width) height, width = spatial_shapes[i] resized_embeddings = F.interpolate( positional_embeddings, size=(height, width), mode="bilinear", align_corners=False, antialias=True, ) # (1, dim, target_height, target_width) -> (target_height * target_width, dim) resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1) # Cast to original dtype resized_embeddings = resized_embeddings.to(source_dtype) resulted_positional_embeddings[i, : height * width] = resized_embeddings resulted_positional_embeddings[i, height * width :] = resized_embeddings[0] return resulted_positional_embeddings def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor: """ Args: pixel_values (`torch.FloatTensor`): Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size) spatial_shapes (`list[tuple[int, int]]`): Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to """ # Apply patch embeddings to already patchified pixel values target_dtype = self.patch_embedding.weight.dtype patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # Get positional resized and padded positional embeddings positional_embeddings = self.position_embedding.weight.reshape( self.position_embedding_size, self.position_embedding_size, -1 ) resized_positional_embeddings = self.resize_positional_embeddings( positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1] ) # Add positional embeddings to patch embeddings embeddings = patch_embeds + resized_positional_embeddings return embeddings def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float = 0.0, **kwargs, ): attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights class Fgclip2Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config): super().__init__() self.config = config self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads if self.head_dim * self.num_heads != self.embed_dim: raise ValueError( f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" f" {self.num_heads})." ) self.scale = self.head_dim**-0.5 self.dropout = config.attention_dropout self.is_causal = False self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" batch_size, seq_length, embed_dim = hidden_states.shape queries = self.q_proj(hidden_states) keys = self.k_proj(hidden_states) values = self.v_proj(hidden_states) queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] attn_output, attn_weights = attention_interface( self, queries, keys, values, attention_mask, is_causal=self.is_causal, scaling=self.scale, dropout=0.0 if not self.training else self.dropout, ) attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() attn_output = self.out_proj(attn_output) return attn_output, attn_weights class Fgclip2MLP(nn.Module): def __init__(self, config): super().__init__() self.config = config self.activation_fn = ACT2FN[config.hidden_act] self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) return hidden_states class Fgclip2EncoderLayer(GradientCheckpointingLayer): def __init__(self, config: Union[Fgclip2VisionConfig, Fgclip2TextConfig]): super().__init__() self.embed_dim = config.hidden_size self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.self_attn = Fgclip2Attention(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = Fgclip2MLP(config) @auto_docstring def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.FloatTensor: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class Fgclip2Encoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Fgclip2EncoderLayer`]. Args: config: Fgclip2Config """ def __init__(self, config: Fgclip2Config): super().__init__() self.config = config self.layers = nn.ModuleList([Fgclip2EncoderLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False # Ignore copy @auto_docstring def forward( self, inputs_embeds, attention_mask: Optional[torch.Tensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutput: hidden_states = inputs_embeds for encoder_layer in self.layers: hidden_states = encoder_layer( hidden_states, attention_mask, **kwargs, ) return BaseModelOutput(last_hidden_state=hidden_states) class Fgclip2VisionTransformer(nn.Module): def __init__(self, config: Fgclip2VisionConfig): super().__init__() self.config = config embed_dim = config.hidden_size self.embeddings = Fgclip2VisionEmbeddings(config) self.encoder = Fgclip2Encoder(config) self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head if self.use_head: self.head = Fgclip2MultiheadAttentionPoolingHead(config) @can_return_tuple @auto_docstring def forward( self, pixel_values: torch.FloatTensor, attention_mask: torch.Tensor, spatial_shapes: torch.LongTensor, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, ) -> BaseModelOutputWithPooling: r""" spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) hidden_states = self.embeddings(pixel_values, spatial_shapes) if attention_mask is not None and self.config._attn_implementation != "flash_attention_2": # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) else: encoder_attention_mask = attention_mask encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, attention_mask=encoder_attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, ) last_hidden_state = encoder_outputs.last_hidden_state last_hidden_state = self.post_layernorm(last_hidden_state) pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooler_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, ) def _trunc_normal_(tensor, mean, std, a, b): # Cut & paste from PyTorch official master until it's in a few official releases - RW # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf def norm_cdf(x): # Computes standard normal cumulative distribution function return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 if (mean < a - 2 * std) or (mean > b + 2 * std): warnings.warn( "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " "The distribution of values may be incorrect.", stacklevel=2, ) # Values are generated by using a truncated uniform distribution and # then using the inverse CDF for the normal distribution. # Get upper and lower cdf values l = norm_cdf((a - mean) / std) u = norm_cdf((b - mean) / std) # Uniformly fill tensor with values from [l, u], then translate to # [2l-1, 2u-1]. tensor.uniform_(2 * l - 1, 2 * u - 1) # Use inverse cdf transform for normal distribution to get truncated # standard normal tensor.erfinv_() # Transform to proper mean, std tensor.mul_(std * math.sqrt(2.0)) tensor.add_(mean) # Clamp to ensure it's in the proper range tensor.clamp_(min=a, max=b) def trunc_normal_tf_( tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0 ) -> torch.Tensor: """Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for generating the random values works best when :math:`a \\leq \text{mean} \\leq b`. NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0 and the result is subsequently scaled and shifted by the mean and std args. Args: tensor: an n-dimensional `torch.Tensor` mean: the mean of the normal distribution std: the standard deviation of the normal distribution a: the minimum cutoff value b: the maximum cutoff value """ with torch.no_grad(): _trunc_normal_(tensor, 0, 1.0, a, b) tensor.mul_(std).add_(mean) def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"): fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) if mode == "fan_in": denom = fan_in elif mode == "fan_out": denom = fan_out elif mode == "fan_avg": denom = (fan_in + fan_out) / 2 variance = scale / denom if distribution == "truncated_normal": # constant is stddev of standard normal truncated to (-2, 2) trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978) elif distribution == "normal": with torch.no_grad(): tensor.normal_(std=math.sqrt(variance)) elif distribution == "uniform": bound = math.sqrt(3 * variance) with torch.no_grad(): tensor.uniform_(-bound, bound) else: raise ValueError(f"invalid distribution {distribution}") def lecun_normal_(tensor): variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal") def default_flax_embed_init(tensor): variance_scaling_(tensor, mode="fan_in", distribution="normal") @auto_docstring class Fgclip2PreTrainedModel(PreTrainedModel): config: Fgclip2Config base_model_prefix = "fgclip2" supports_gradient_checkpointing = True _no_split_modules = [ "Fgclip2TextEmbeddings", "Fgclip2VisionEmbeddings", "Fgclip2EncoderLayer", "Fgclip2MultiheadAttentionPoolingHead", ] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": Fgclip2EncoderLayer, "attentions": Fgclip2Attention, } def _init_weights(self, module): """Initialize the weights""" if isinstance(module, Fgclip2VisionEmbeddings): width = ( self.config.vision_config.hidden_size if isinstance(self.config, Fgclip2Config) else self.config.hidden_size ) nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width)) elif isinstance(module, nn.Embedding): default_flax_embed_init(module.weight) elif isinstance(module, Fgclip2Attention): nn.init.xavier_uniform_(module.q_proj.weight) nn.init.xavier_uniform_(module.k_proj.weight) nn.init.xavier_uniform_(module.v_proj.weight) nn.init.xavier_uniform_(module.out_proj.weight) nn.init.zeros_(module.q_proj.bias) nn.init.zeros_(module.k_proj.bias) nn.init.zeros_(module.v_proj.bias) nn.init.zeros_(module.out_proj.bias) elif isinstance(module, Fgclip2MLP): nn.init.xavier_uniform_(module.fc1.weight) nn.init.xavier_uniform_(module.fc2.weight) nn.init.normal_(module.fc1.bias, std=1e-6) nn.init.normal_(module.fc2.bias, std=1e-6) elif isinstance(module, Fgclip2MultiheadAttentionPoolingHead): nn.init.xavier_uniform_(module.probe.data) nn.init.xavier_uniform_(module.attention.in_proj_weight.data) nn.init.zeros_(module.attention.in_proj_bias.data) elif isinstance(module, Fgclip2Model): logit_scale_init = torch.log(torch.tensor(1.0)) module.logit_scale.data.fill_(logit_scale_init) module.logit_bias.data.zero_() elif isinstance(module, (nn.Linear, nn.Conv2d)): lecun_normal_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) class Fgclip2TextEmbeddings(nn.Module): def __init__(self, config: Fgclip2TextConfig): super().__init__() embed_dim = config.hidden_size self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) # position_ids (1, len position emb) is contiguous in memory and exported when serialized self.register_buffer( "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False ) keep_len = config.keep_len longtext_len = config.longtext_len self.position_embedding_res = nn.Embedding(longtext_len, embed_dim) self.position_embedding_ori = nn.Embedding(longtext_len, embed_dim) self.mask1 = torch.zeros([longtext_len, 1]) self.mask1[:keep_len, :] = 1 self.mask2 = torch.zeros([longtext_len, 1]) self.mask2[keep_len:, :] = 1 # position_ids (1, len position emb) is contiguous in memory and exported when serialized self.register_buffer("position_ids", torch.arange(longtext_len).expand((1, -1)), persistent=False) def forward( self, input_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_short_position_ids: Optional[bool] = True, ) -> torch.Tensor: r""" Args: use_short_position_ids (`bool`, optional, defaults to `True`): If `True`, applies a positional encoding scheme optimized for **short-text processing** and **local-region description processing**, such as phrases or simple sentences. Corresponds to the `"short"` and `"box"` walk type. Assumes compact semantic structure and local dependency dominance. """ seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] if position_ids is None: position_ids = self.position_ids[:, :seq_length] if inputs_embeds is None: inputs_embeds = self.token_embedding(input_ids) if use_short_position_ids: position_embeddings = self.position_embedding(position_ids) embeddings = inputs_embeds + position_embeddings else: position_embeddings_res = self.position_embedding_res(position_ids) position_embeddings_ori = self.position_embedding_ori(position_ids) embeddings = ( inputs_embeds + (position_embeddings_ori * self.mask1.to(inputs_embeds.device)) .type(inputs_embeds.dtype) .to(inputs_embeds.device) + (position_embeddings_res * self.mask2.to(inputs_embeds.device)) .type(inputs_embeds.dtype) .to(inputs_embeds.device) ) return embeddings class Fgclip2TextTransformer(nn.Module): def __init__(self, config: Fgclip2TextConfig): super().__init__() self.config = config embed_dim = config.hidden_size self.embeddings = Fgclip2TextEmbeddings(config) self.encoder = Fgclip2Encoder(config) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) self.head = nn.Linear(embed_dim, config.projection_size) @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, walk_type: str = "short", # Modified: Single parameter **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" Args: walk_type (`str`, optional, defaults to `"short"`): The traversal strategy used during feature extraction. Must be one of `"short"`, `"box"`, or `"long"`. This controls how contextual information is aggregated across the input: - `"short"`: Optimized for short-text understanding, focusing on tight semantic coherence and direct word interactions. Suitable when the input is a phrase or brief sentence. - `"box"`: Designed for local-region description processing, such as grounding in vision-language models or processing localized textual descriptions (e.g., object regions or segments). Emphasizes dense features within bounded semantic units. - `"long"`: Tailored for long-form text processing, enabling modeling of extended dependencies and discourse structure. Uses strategies like chunking or hierarchical attention to handle longer sequences effectively. """ if input_ids is None: raise ValueError("You have to specify input_ids") # Validate walk_type walk_type = walk_type.lower() if walk_type not in ["short", "box", "long"]: raise ValueError(f"Invalid `walk_type`: {walk_type}. Must be one of 'short', 'box', 'long'.") # Convert walk_type to boolean flags for internal logic walk_short = walk_type == "short" walk_box = walk_type == "box" walk_long = walk_type == "long" input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) hidden_states = self.embeddings( input_ids=input_ids, position_ids=position_ids, use_short_position_ids=(not walk_long) ) # note: fgclip2's text model does not use a causal mask, unlike the original CLIP model. # expand attention_mask uses_flash_attention = "flash" in self.config._attn_implementation if uses_flash_attention: attention_mask = None elif attention_mask is not None and not uses_flash_attention: # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len] attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, **kwargs, ) last_hidden_state = encoder_outputs.last_hidden_state last_hidden_state = self.final_layer_norm(last_hidden_state) # The model uses the last token's hidden state, which may be padding. pooled_output = last_hidden_state[:, -1, :] if walk_short == True: assert walk_box == False assert walk_long == False temp_pool_out = [] for i in range(pooled_output.shape[0]): temp_pool_out.append(self.head(pooled_output[i : i + 1])) pooled_output = torch.cat(temp_pool_out, dim=0) # pooled_output = self.head(pooled_output) if walk_box == True: assert walk_short == False assert walk_long == False pooled_output = pooled_output if walk_long == True: assert walk_short == False assert walk_box == False pooled_output = pooled_output return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooled_output, ) @auto_docstring( custom_intro=""" The text model from Fgclip2 without any head or projection on top. """ ) class Fgclip2TextModel(Fgclip2PreTrainedModel): config: Fgclip2TextConfig def __init__(self, config: Fgclip2TextConfig): super().__init__(config) self.text_model = Fgclip2TextTransformer(config) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> nn.Module: return self.text_model.embeddings.token_embedding def set_input_embeddings(self, value): self.text_model.embeddings.token_embedding = value @check_model_inputs @auto_docstring def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, walk_type: str = "short", # Modified: Single parameter **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" Args: walk_type (`str`, optional, defaults to `"short"`): The traversal strategy used during feature extraction. Must be one of `"short"`, `"box"`, or `"long"`. This controls how contextual information is aggregated across the input: - `"short"`: Optimized for short-text understanding, focusing on tight semantic coherence and direct word interactions. Suitable when the input is a phrase or brief sentence. - `"box"`: Designed for local-region description processing, such as grounding in vision-language models or processing localized textual descriptions (e.g., object regions or segments). Emphasizes dense features within bounded semantic units. - `"long"`: Tailored for long-form text processing, enabling modeling of extended dependencies and discourse structure. Uses strategies like chunking or hierarchical attention to handle longer sequences effectively. """ return self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, walk_type=walk_type, # Modified: Pass single parameter **kwargs, ) class Fgclip2MultiheadAttentionPoolingHead(nn.Module): """Multihead Attention Pooling.""" def __init__(self, config: Fgclip2VisionConfig): super().__init__() self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.mlp = Fgclip2MLP(config) self.num_heads = config.num_attention_heads def forward(self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: batch_size = hidden_state.shape[0] probe = self.probe.repeat(batch_size, 1, 1) if attention_mask is not None: target_len, source_len = probe.shape[1], hidden_state.shape[1] attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len) attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1) attention_mask = attention_mask.reshape(-1, target_len, source_len) group_size = self.num_heads outputs = [] for i in range(batch_size): start_idx = i * group_size end_idx = start_idx + group_size out_i = self.attention( probe[i : i + 1], hidden_state[i : i + 1], hidden_state[i : i + 1], attn_mask=attention_mask[start_idx:end_idx] if attention_mask is not None else None, )[0] outputs.append(out_i) hidden_state = torch.cat(outputs, dim=0) residual = hidden_state hidden_state = self.layernorm(hidden_state) temp_outs = [] for k in range(batch_size): out_k = self.mlp(hidden_state[k : k + 1]) temp_outs.append(out_k) hidden_state = residual + torch.cat(temp_outs, dim=0) return hidden_state[:, 0] @auto_docstring( custom_intro=""" The vision model from Fgclip2 without any head or projection on top. """ ) class Fgclip2VisionModel(Fgclip2PreTrainedModel): config: Fgclip2VisionConfig main_input_name = "pixel_values" def __init__(self, config: Fgclip2VisionConfig): super().__init__(config) self.vision_model = Fgclip2VisionTransformer(config) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> nn.Module: return self.vision_model.embeddings.patch_embedding @check_model_inputs @auto_docstring def forward( self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.Tensor, spatial_shapes: torch.LongTensor, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, ) -> BaseModelOutputWithPooling: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, Fgclip2VisionModel >>> model = Fgclip2VisionModel.from_pretrained("qihoo360/fg-clip2-base") >>> processor = AutoProcessor.from_pretrained("qihoo360/fg-clip2-base") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooled_output = outputs.pooler_output # pooled features ```""" return self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, output_attentions=output_attentions, output_hidden_states=output_hidden_states, ) @auto_docstring class Fgclip2Model(Fgclip2PreTrainedModel): config: Fgclip2Config def __init__(self, config: Fgclip2Config): super().__init__(config) if not isinstance(config.text_config, Fgclip2TextConfig): raise TypeError( "config.text_config is expected to be of type Fgclip2TextConfig but is of type" f" {type(config.text_config)}." ) if not isinstance(config.vision_config, Fgclip2VisionConfig): raise TypeError( "config.vision_config is expected to be of type Fgclip2VisionConfig but is of type" f" {type(config.vision_config)}." ) text_config = config.text_config vision_config = config.vision_config # First, initialize the text and vision models with proper attention implementation text_model = Fgclip2TextModel._from_config(text_config) vision_model = Fgclip2VisionModel._from_config(vision_config) # Second, get the text and vision submodules (for backward compatibility) self.text_model = text_model.text_model self.vision_model = vision_model.vision_model self.logit_scale = nn.Parameter(torch.randn(1)) self.logit_bias = nn.Parameter(torch.randn(1)) self.dense_feature_head = Fgclip2MultiheadAttentionPoolingHead(vision_config) self.embed_dim = text_config.hidden_size self.longtext_head = nn.Linear(self.embed_dim, self.embed_dim) self.boxtext_head = nn.Linear(self.embed_dim, self.embed_dim) # Initialize weights and apply final processing self.post_init() @filter_out_non_signature_kwargs() @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, walk_type: str = "short", ) -> torch.FloatTensor: r""" Extracts feature representations from the input text. Args: input_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`): The token IDs of the input sequence, as generated by the tokenizer. attention_mask (`torch.Tensor`, optional, of shape `(batch_size, sequence_length)`): A mask indicating which tokens are valid (1) and which are padding (0). If not provided, all tokens are assumed to be valid. position_ids (`torch.Tensor`, optional, of shape `(batch_size, sequence_length)`): Position indices for each token in the sequence. If not provided, positions are automatically constructed based on `input_ids`. walk_type (`str`, optional, defaults to `"short"`): The traversal strategy used during feature extraction. Must be one of `"short"`, `"box"`, or `"long"`. This controls how contextual information is aggregated across the input: - `"short"`: Optimized for short-text understanding, focusing on tight semantic coherence and direct word interactions. Suitable when the input is a phrase or brief sentence. - `"box"`: Designed for local-region description processing, such as grounding in vision-language models or processing localized textual descriptions (e.g., object regions or segments). Emphasizes dense features within bounded semantic units. - `"long"`: Tailored for long-form text processing, enabling modeling of extended dependencies and discourse structure. Uses strategies like chunking or hierarchical attention to handle longer sequences effectively. Returns: `torch.FloatTensor` of shape `(batch_size, hidden_size)` or `(batch_size, sequence_length, hidden_size)`: The extracted feature tensor representing the input text. The output shape depends on whether a pooled representation or per-token embeddings are returned. Examples: ```python >>> from transformers import AutoTokenizer, AutoModel >>> import torch >>> model = AutoModel.from_pretrained("qihoo360/fg-clip2-base") >>> tokenizer = AutoTokenizer.from_pretrained("qihoo360/fg-clip2-base") >>> # important: make sure to set padding="max_length" as that's how the model was trained >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt") >>> with torch.no_grad(): ... text_features = model.get_text_features(**inputs, walk_type="short") ```""" walk_type = walk_type.lower() if walk_type not in ["short", "box", "long"]: raise ValueError(f"Invalid `walk_type`: {walk_type}. Must be one of 'short', 'box', 'long'.") walk_short = walk_type == "short" walk_box = walk_type == "box" walk_long = walk_type == "long" text_outputs: BaseModelOutputWithPooling = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, walk_type=walk_type, ) if walk_short: pooled_output = text_outputs.pooler_output if walk_box: pooled_output = self.boxtext_head(text_outputs.pooler_output) if walk_long: pooled_output = self.longtext_head(text_outputs.pooler_output) return pooled_output @filter_out_non_signature_kwargs() @auto_docstring def get_image_features( self, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.Tensor] = None, spatial_shapes: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`Fgclip2VisionModel`]. Examples: ```python >>> import torch >>> from transformers import AutoProcessor, AutoModel >>> from transformers.image_utils import load_image >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = load_image(url) >>> model = AutoModel.from_pretrained("qihoo360/fg-clip2-base") >>> processor = AutoProcessor.from_pretrained("qihoo360/fg-clip2-base") >>> inputs = processor(images=image, return_tensors="pt") >>> with torch.no_grad(): ... image_features = model.get_image_features(**inputs) ``` """ vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, ) pooled_output = vision_outputs.pooler_output return pooled_output # NOTE: Fgclip2Model uses Pretrained backbones, so we don't need to add `check_model_inputs` here @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.Tensor] = None, spatial_shapes: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, return_loss: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, walk_type: str = "short", ) -> Fgclip2Output: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. return_loss (`bool`, *optional*): Whether or not to return the contrastive loss. walk_type (`str`, optional, defaults to `"short"`): The traversal strategy used during feature extraction. Must be one of `"short"`, `"box"`, or `"long"`. This controls how contextual information is aggregated across the input: - `"short"`: Optimized for short-text understanding, focusing on tight semantic coherence and direct word interactions. Suitable when the input is a phrase or brief sentence. - `"box"`: Designed for local-region description processing, such as grounding in vision-language models or processing localized textual descriptions (e.g., object regions or segments). Emphasizes dense features within bounded semantic units. - `"long"`: Tailored for long-form text processing, enabling modeling of extended dependencies and discourse structure. Uses strategies like chunking or hierarchical attention to handle longer sequences effectively. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, AutoModel >>> import torch >>> model = AutoModel.from_pretrained("qihoo360/fg-clip2-base") >>> processor = AutoProcessor.from_pretrained("qihoo360/fg-clip2-base") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> texts = ["a photo of 2 cats", "a photo of 2 dogs"] >>> # important: we pass `padding=max_length` since the model was trained with this >>> inputs = processor(text=texts, images=image, padding="max_length", return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**inputs) >>> logits_per_image = outputs.logits_per_image >>> probs = torch.sigmoid(logits_per_image) # these are the probabilities >>> print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'") 31.9% that image 0 is 'a photo of 2 cats' ``` """ walk_type = walk_type.lower() if walk_type not in ["short", "box", "long"]: raise ValueError(f"Invalid `walk_type`: {walk_type}. Must be one of 'short', 'box', 'long'.") walk_short = walk_type == "short" walk_box = walk_type == "box" walk_long = walk_type == "long" # Use Fgclip2 model's config for some fields (if specified) instead of those of vision & text components. output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, output_attentions=output_attentions, output_hidden_states=output_hidden_states, ) text_outputs: BaseModelOutputWithPooling = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, walk_type=walk_type, ) image_embeds = vision_outputs.pooler_output if walk_short: text_embeds = text_outputs.pooler_output if walk_box: text_embeds = self.boxtext_head(text_outputs.pooler_output) if walk_long: text_embeds = self.longtext_head(text_outputs.pooler_output) # normalized features image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) # cosine similarity as logits logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) logit_scale, logit_bias = self.logit_scale.to(text_embeds.device), self.logit_bias.to(text_embeds.device) logits_per_text = logits_per_text * logit_scale.exp() + logit_bias logits_per_image = logits_per_text.t() loss = None if return_loss: # Adapted from https://github.com/google-research/big_vision/blob/01edb81a4716f93a48be43b3a4af14e29cdb3a7f/big_vision/trainers/proj/image_text/fgclip2.py#L287 eye = torch.eye(logits_per_text.size(0), device=logits_per_text.device) m1_diag1 = -torch.ones_like(logits_per_text) + 2 * eye loglik = torch.nn.functional.logsigmoid(m1_diag1 * logits_per_text) nll = -torch.sum(loglik, dim=-1) loss = nll.mean() return Fgclip2Output( loss=loss, logits_per_image=logits_per_image, logits_per_text=logits_per_text, text_embeds=text_embeds, image_embeds=image_embeds, text_model_output=text_outputs, vision_model_output=vision_outputs, ) # New function: Acquire dense visual features of images with support for dynamic resolution @filter_out_non_signature_kwargs() @auto_docstring def get_image_dense_feature( self, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.Tensor] = None, spatial_shapes: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: r""" Extract dense visual features from input images by forwarding through the vision backbone. Args: pixel_values (`torch.FloatTensor`): Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size) pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. Returns: `torch.FloatTensor` of shape `(batch_size, max_num_patches, hidden_size)`: """ vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, ) probe = vision_outputs.last_hidden_state hidden_state = vision_outputs.last_hidden_state attention_mask = pixel_attention_mask if attention_mask is not None: target_len, source_len = probe.shape[1], hidden_state.shape[1] attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len) attention_mask = attention_mask.repeat(1, self.dense_feature_head.num_heads, 1, 1) attention_mask = attention_mask.reshape(-1, target_len, source_len) hidden_state = self.dense_feature_head.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[ 0 ] residual = hidden_state hidden_state = self.dense_feature_head.layernorm(hidden_state) hidden_state = residual + self.dense_feature_head.mlp(hidden_state) feature_map = hidden_state return feature_map # New function: Acquire local features of images, applicable to retrieval, classification, and localization, with support for dynamic resolution @filter_out_non_signature_kwargs() @auto_docstring def get_image_region_features( self, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.Tensor] = None, spatial_shapes: Optional[torch.LongTensor] = None, image_sizes: Optional[list[tuple]] = None, region_infos: Optional[list[list[list[float]]]] = None, ) -> list[torch.FloatTensor]: r""" Extract region-of-interest (RoI) features from images using RoI Align. This method supports batched processing of variable-sized images and allows feature extraction from user-specified image regions. The input can be either a full image with corresponding region coordinates. Features are extracted per region (e.g., bounding boxes), making this function suitable for tasks such as object detection, referring expression grounding, or vision-language alignment. Args: pixel_values (`torch.FloatTensor`): Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size) pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. image_sizes (`List[tuple]`, optional, each tuple of form `(int, int)`): Original size (height, width) of each image in the batch before padding or resizing. Required for accurate coordinate projection when region_infos are defined in original image space. region_infos (`List[List[List[float]]]`, optional): Bounding box coordinates for regions of interest in each image. Format: - Outer list: length `batch_size` - Middle list: number of regions per image - Inner list: each contains `[x_min, y_min, x_max, y_max]` in **absolute pixel coordinates** relative to the original image size (as specified in `image_sizes`). These boxes are projected to feature map space using `image_sizes` and `spatial_shapes`, then used to pool features via RoI Align or equivalent. Returns: `List[torch.FloatTensor]`: A list of length `batch_size`, where each element is a tensor of shape `(num_boxes, hidden_dim)` containing the extracted visual features for each region in the corresponding image. Example:: >>> # For a batch of 2 images >>> region_features = model.get_image_region_features( >>> pixel_values=pixel_values, >>> image_sizes=[(640, 480), (480, 640)], >>> region_infos=[ >>> [[100, 100, 200, 200], [300, 300, 400, 400]], # 2 boxes in first image >>> [[50, 50, 150, 150]] # 1 box in second image >>> ] >>> ) >>> print(region_features[0].shape) # torch.Size([2, hidden_dim]) >>> print(region_features[1].shape) # torch.Size([1, hidden_dim]) """ if region_infos is None or len(region_infos) == 0: return [] # Get dense feature maps: (B, N, D) dense_feature_map = self.get_image_dense_feature( pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, ) bs, _, hidden_dim = dense_feature_map.shape all_region_features = [] for i in range(bs): h, w = spatial_shapes[i].tolist() img_h, img_w = image_sizes[i] bboxes = region_infos[i] if not bboxes: all_region_features.append(torch.empty(0, hidden_dim, device=dense_feature_map.device)) continue # Reshape to (1, C, H', W') num_valid = h * w feat_seq = dense_feature_map[i, :num_valid] # (num_valid, D) feat_map = feat_seq.view(h, w, hidden_dim).permute(2, 0, 1).unsqueeze(0) # (1, D, H', W') # Normalize bboxes to feature map coordinates rois = [] for x1, y1, x2, y2 in bboxes: nx1 = (x1 / img_w) * w ny1 = (y1 / img_h) * h nx2 = (x2 / img_w) * w ny2 = (y2 / img_h) * h rois.append([0, nx1, ny1, nx2, ny2]) # rois_tensor = torch.tensor(rois, dtype=torch.float32, device=feat_map.device) # (N, 5) # RoI Align on single image pooled = roi_align( input=feat_map, boxes=rois_tensor, output_size=(1, 1), spatial_scale=1.0, sampling_ratio=-1, aligned=True, ) # (N, D, 1, 1) region_feats = pooled.squeeze(-1).squeeze(-1) # (N, D) all_region_features.append(region_feats) return all_region_features __all__ = ["Fgclip2Model", "Fgclip2PreTrainedModel", "Fgclip2TextModel", "Fgclip2VisionModel"]