Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 Google AI and The HuggingFace Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch OWLv2 model."""16 17from dataclasses import dataclass18from functools import lru_cache19from typing import Any, Optional, Union20 21import torch22from torch import Tensor, nn23 24from ...activations import ACT2FN25from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask26from ...modeling_layers import GradientCheckpointingLayer27from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling28from ...modeling_utils import PreTrainedModel29from ...utils import (30 ModelOutput,31 auto_docstring,32 filter_out_non_signature_kwargs,33 is_vision_available,34 logging,35 torch_int,36)37from .configuration_owlv2 import Owlv2Config, Owlv2TextConfig, Owlv2VisionConfig38 39 40if is_vision_available():41 from transformers.image_transforms import center_to_corners_format42 43 44logger = logging.get_logger(__name__)45 46 47# See all Owlv2 models at https://huggingface.co/models?filter=owlv248 49 50# Copied from transformers.models.clip.modeling_clip.contrastive_loss with clip->owlv251def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:52 return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))53 54 55# Copied from transformers.models.clip.modeling_clip.clip_loss with clip->owlv256def owlv2_loss(similarity: torch.Tensor) -> torch.Tensor:57 caption_loss = contrastive_loss(similarity)58 image_loss = contrastive_loss(similarity.t())59 return (caption_loss + image_loss) / 2.060 61 62@dataclass63@auto_docstring64class Owlv2Output(ModelOutput):65 r"""66 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):67 Contrastive loss for image-text similarity.68 logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):69 The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text70 similarity scores.71 logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):72 The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image73 similarity scores.74 text_embeds (`torch.FloatTensor` of shape `(batch_size * num_max_text_queries, output_dim`):75 The text embeddings obtained by applying the projection layer to the pooled output of [`Owlv2TextModel`].76 image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):77 The image embeddings obtained by applying the projection layer to the pooled output of78 [`Owlv2VisionModel`].79 text_model_output (tuple[`BaseModelOutputWithPooling`]):80 The output of the [`Owlv2TextModel`].81 vision_model_output (`BaseModelOutputWithPooling`):82 The output of the [`Owlv2VisionModel`].83 """84 85 loss: Optional[torch.FloatTensor] = None86 logits_per_image: Optional[torch.FloatTensor] = None87 logits_per_text: Optional[torch.FloatTensor] = None88 text_embeds: Optional[torch.FloatTensor] = None89 image_embeds: Optional[torch.FloatTensor] = None90 text_model_output: BaseModelOutputWithPooling = None91 vision_model_output: BaseModelOutputWithPooling = None92 93 def to_tuple(self) -> tuple[Any]:94 return tuple(95 self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()96 for k in self.keys()97 )98 99 100# Copied from transformers.loss.loss_for_object_detection._upcast101def _upcast(t: Tensor) -> Tensor:102 # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type103 if t.is_floating_point():104 return t if t.dtype in (torch.float32, torch.float64) else t.float()105 else:106 return t if t.dtype in (torch.int32, torch.int64) else t.int()107 108 109# Copied from transformers.loss.loss_for_object_detection.box_area110def box_area(boxes: Tensor) -> Tensor:111 """112 Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.113 114 Args:115 boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):116 Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1117 < x2` and `0 <= y1 < y2`.118 119 Returns:120 `torch.FloatTensor`: a tensor containing the area for each box.121 """122 boxes = _upcast(boxes)123 return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])124 125 126# Copied from transformers.loss.loss_for_object_detection.box_iou127def box_iou(boxes1, boxes2):128 area1 = box_area(boxes1)129 area2 = box_area(boxes2)130 131 left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]132 right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]133 134 width_height = (right_bottom - left_top).clamp(min=0) # [N,M,2]135 inter = width_height[:, :, 0] * width_height[:, :, 1] # [N,M]136 137 union = area1[:, None] + area2 - inter138 139 iou = inter / union140 return iou, union141 142 143# Copied from transformers.loss.loss_for_object_detection.generalized_box_iou144def generalized_box_iou(boxes1, boxes2):145 """146 Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.147 148 Returns:149 `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)150 """151 # degenerate boxes gives inf / nan results152 # so do an early check153 if not (boxes1[:, 2:] >= boxes1[:, :2]).all():154 raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")155 if not (boxes2[:, 2:] >= boxes2[:, :2]).all():156 raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")157 iou, union = box_iou(boxes1, boxes2)158 159 top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])160 bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])161 162 width_height = (bottom_right - top_left).clamp(min=0) # [N,M,2]163 area = width_height[:, :, 0] * width_height[:, :, 1]164 165 return iou - (area - union) / area166 167 168@dataclass169@auto_docstring(170 custom_intro="""171 Output type of [`Owlv2ForObjectDetection`].172 """173)174class Owlv2ObjectDetectionOutput(ModelOutput):175 r"""176 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):177 Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a178 bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized179 scale-invariant IoU loss.180 loss_dict (`Dict`, *optional*):181 A dictionary containing the individual losses. Useful for logging.182 logits (`torch.FloatTensor` of shape `(batch_size, num_patches, num_queries)`):183 Classification logits (including no-object) for all queries.184 objectness_logits (`torch.FloatTensor` of shape `(batch_size, num_patches, 1)`):185 The objectness logits of all image patches. OWL-ViT represents images as a set of image patches where the186 total number of patches is (image_size / patch_size)**2.187 pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 4)`):188 Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These189 values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding190 possible padding). You can use [`~Owlv2ImageProcessor.post_process_object_detection`] to retrieve the191 unnormalized bounding boxes.192 text_embeds (`torch.FloatTensor` of shape `(batch_size, num_max_text_queries, output_dim`):193 The text embeddings obtained by applying the projection layer to the pooled output of [`Owlv2TextModel`].194 image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):195 Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes image196 embeddings for each patch.197 class_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`):198 Class embeddings of all image patches. OWLv2 represents images as a set of image patches where the total199 number of patches is (image_size / patch_size)**2.200 text_model_output (tuple[`BaseModelOutputWithPooling`]):201 The output of the [`Owlv2TextModel`].202 vision_model_output (`BaseModelOutputWithPooling`):203 The output of the [`Owlv2VisionModel`].204 """205 206 loss: Optional[torch.FloatTensor] = None207 loss_dict: Optional[dict] = None208 logits: Optional[torch.FloatTensor] = None209 objectness_logits: Optional[torch.FloatTensor] = None210 pred_boxes: Optional[torch.FloatTensor] = None211 text_embeds: Optional[torch.FloatTensor] = None212 image_embeds: Optional[torch.FloatTensor] = None213 class_embeds: Optional[torch.FloatTensor] = None214 text_model_output: BaseModelOutputWithPooling = None215 vision_model_output: BaseModelOutputWithPooling = None216 217 def to_tuple(self) -> tuple[Any]:218 return tuple(219 self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()220 for k in self.keys()221 )222 223 224@dataclass225@auto_docstring(226 custom_intro="""227 Output type of [`Owlv2ForObjectDetection.image_guided_detection`].228 """229)230# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTImageGuidedObjectDetectionOutput with OwlViT->Owlv2,OWL-ViT->OWLv2231class Owlv2ImageGuidedObjectDetectionOutput(ModelOutput):232 r"""233 logits (`torch.FloatTensor` of shape `(batch_size, num_patches, num_queries)`):234 Classification logits (including no-object) for all queries.235 image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):236 Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes237 image embeddings for each patch.238 query_image_embeds (`torch.FloatTensor` of shape `(batch_size, patch_size, patch_size, output_dim`):239 Pooled output of [`Owlv2VisionModel`]. OWLv2 represents images as a set of image patches and computes240 image embeddings for each patch.241 target_pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 4)`):242 Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These243 values are normalized in [0, 1], relative to the size of each individual target image in the batch244 (disregarding possible padding). You can use [`~Owlv2ImageProcessor.post_process_object_detection`] to245 retrieve the unnormalized bounding boxes.246 query_pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_patches, 4)`):247 Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These248 values are normalized in [0, 1], relative to the size of each individual query image in the batch249 (disregarding possible padding). You can use [`~Owlv2ImageProcessor.post_process_object_detection`] to250 retrieve the unnormalized bounding boxes.251 class_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`):252 Class embeddings of all image patches. OWLv2 represents images as a set of image patches where the total253 number of patches is (image_size / patch_size)**2.254 text_model_output (tuple[`BaseModelOutputWithPooling`]):255 The output of the [`Owlv2TextModel`].256 vision_model_output (`BaseModelOutputWithPooling`):257 The output of the [`Owlv2VisionModel`].258 """259 260 logits: Optional[torch.FloatTensor] = None261 image_embeds: Optional[torch.FloatTensor] = None262 query_image_embeds: Optional[torch.FloatTensor] = None263 target_pred_boxes: Optional[torch.FloatTensor] = None264 query_pred_boxes: Optional[torch.FloatTensor] = None265 class_embeds: Optional[torch.FloatTensor] = None266 text_model_output: BaseModelOutputWithPooling = None267 vision_model_output: BaseModelOutputWithPooling = None268 269 def to_tuple(self) -> tuple[Any]:270 return tuple(271 self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()272 for k in self.keys()273 )274 275 276# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionEmbeddings with OwlViT->Owlv2277class Owlv2VisionEmbeddings(nn.Module):278 def __init__(self, config: Owlv2VisionConfig):279 super().__init__()280 self.patch_size = config.patch_size281 self.config = config282 self.embed_dim = config.hidden_size283 self.class_embedding = nn.Parameter(torch.randn(config.hidden_size))284 285 self.patch_embedding = nn.Conv2d(286 in_channels=config.num_channels,287 out_channels=self.embed_dim,288 kernel_size=config.patch_size,289 stride=config.patch_size,290 bias=False,291 )292 293 self.num_patches = (config.image_size // config.patch_size) ** 2294 self.num_positions = self.num_patches + 1295 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)296 self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)297 298 # Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings.interpolate_pos_encoding299 def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:300 """301 This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution302 images. This method is also adapted to support torch.jit tracing.303 304 Adapted from:305 - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and306 - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211307 """308 309 num_patches = embeddings.shape[1] - 1310 position_embedding = self.position_embedding.weight.unsqueeze(0)311 num_positions = position_embedding.shape[1] - 1312 313 # always interpolate when tracing to ensure the exported model works for dynamic input shapes314 if not torch.jit.is_tracing() and num_patches == num_positions and height == width:315 return self.position_embedding(self.position_ids)316 317 class_pos_embed = position_embedding[:, :1]318 patch_pos_embed = position_embedding[:, 1:]319 320 dim = embeddings.shape[-1]321 322 new_height = height // self.patch_size323 new_width = width // self.patch_size324 325 sqrt_num_positions = torch_int(num_positions**0.5)326 patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)327 patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)328 329 patch_pos_embed = nn.functional.interpolate(330 patch_pos_embed,331 size=(new_height, new_width),332 mode="bicubic",333 align_corners=False,334 )335 336 patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)337 338 return torch.cat((class_pos_embed, patch_pos_embed), dim=1)339 340 def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:341 batch_size, _, height, width = pixel_values.shape342 patch_embeds = self.patch_embedding(pixel_values) # shape = [batch_size, num_channels, height, width]343 patch_embeds = patch_embeds.flatten(2).transpose(1, 2)344 345 class_embeds = self.class_embedding.expand(batch_size, 1, -1)346 embeddings = torch.cat([class_embeds, patch_embeds], dim=1)347 if interpolate_pos_encoding:348 embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)349 else:350 embeddings = embeddings + self.position_embedding(self.position_ids)351 return embeddings352 353 354# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextEmbeddings with OwlViT->Owlv2355class Owlv2TextEmbeddings(nn.Module):356 def __init__(self, config: Owlv2TextConfig):357 super().__init__()358 self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)359 self.position_embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)360 361 # position_ids (1, len position emb) is contiguous in memory and exported when serialized362 self.register_buffer(363 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False364 )365 366 def forward(367 self,368 input_ids: Optional[torch.LongTensor] = None,369 position_ids: Optional[torch.LongTensor] = None,370 inputs_embeds: Optional[torch.FloatTensor] = None,371 ) -> torch.Tensor:372 seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]373 374 if position_ids is None:375 position_ids = self.position_ids[:, :seq_length]376 377 if inputs_embeds is None:378 inputs_embeds = self.token_embedding(input_ids)379 380 position_embeddings = self.position_embedding(position_ids)381 embeddings = inputs_embeds + position_embeddings382 383 return embeddings384 385 386# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTAttention with OwlViT->Owlv2387class Owlv2Attention(nn.Module):388 """Multi-headed attention from 'Attention Is All You Need' paper"""389 390 def __init__(self, config):391 super().__init__()392 self.config = config393 self.embed_dim = config.hidden_size394 self.num_heads = config.num_attention_heads395 self.head_dim = self.embed_dim // self.num_heads396 if self.head_dim * self.num_heads != self.embed_dim:397 raise ValueError(398 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"399 f" {self.num_heads})."400 )401 self.scale = self.head_dim**-0.5402 self.dropout = config.attention_dropout403 404 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)405 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)406 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)407 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)408 409 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):410 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()411 412 def forward(413 self,414 hidden_states: torch.Tensor,415 attention_mask: Optional[torch.Tensor] = None,416 causal_attention_mask: Optional[torch.Tensor] = None,417 output_attentions: Optional[bool] = False,418 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:419 """Input shape: Batch x Time x Channel"""420 421 bsz, tgt_len, embed_dim = hidden_states.size()422 423 # get query proj424 query_states = self.q_proj(hidden_states) * self.scale425 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)426 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)427 428 proj_shape = (bsz * self.num_heads, -1, self.head_dim)429 query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)430 key_states = key_states.view(*proj_shape)431 value_states = value_states.view(*proj_shape)432 433 src_len = key_states.size(1)434 attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))435 436 if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):437 raise ValueError(438 f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"439 f" {attn_weights.size()}"440 )441 442 # apply the causal_attention_mask first443 if causal_attention_mask is not None:444 if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):445 raise ValueError(446 f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"447 f" {causal_attention_mask.size()}"448 )449 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask450 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)451 452 if attention_mask is not None:453 if attention_mask.size() != (bsz, 1, tgt_len, src_len):454 raise ValueError(455 f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"456 )457 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask458 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)459 460 attn_weights = nn.functional.softmax(attn_weights, dim=-1)461 462 if output_attentions:463 # this operation is a bit awkward, but it's required to464 # make sure that attn_weights keeps its gradient.465 # In order to do so, attn_weights have to reshaped466 # twice and have to be reused in the following467 attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)468 attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)469 else:470 attn_weights_reshaped = None471 472 attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)473 474 # For int8 compatibility, sometimes the `attn_probs` are in `fp32`475 attn_probs = attn_probs.to(value_states.dtype)476 477 attn_output = torch.bmm(attn_probs, value_states)478 479 if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):480 raise ValueError(481 f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"482 f" {attn_output.size()}"483 )484 485 attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)486 attn_output = attn_output.transpose(1, 2)487 attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)488 489 attn_output = self.out_proj(attn_output)490 491 return attn_output, attn_weights_reshaped492 493 494# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Owlv2495class Owlv2MLP(nn.Module):496 def __init__(self, config):497 super().__init__()498 self.config = config499 self.activation_fn = ACT2FN[config.hidden_act]500 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)501 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)502 503 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:504 hidden_states = self.fc1(hidden_states)505 hidden_states = self.activation_fn(hidden_states)506 hidden_states = self.fc2(hidden_states)507 return hidden_states508 509 510# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->Owlv2511class Owlv2EncoderLayer(GradientCheckpointingLayer):512 def __init__(self, config: Owlv2Config):513 super().__init__()514 self.embed_dim = config.hidden_size515 self.self_attn = Owlv2Attention(config)516 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)517 self.mlp = Owlv2MLP(config)518 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)519 520 def forward(521 self,522 hidden_states: torch.Tensor,523 attention_mask: torch.Tensor,524 causal_attention_mask: torch.Tensor,525 output_attentions: Optional[bool] = False,526 ) -> tuple[torch.FloatTensor]:527 """528 Args:529 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`530 attention_mask (`torch.FloatTensor`): attention mask of size531 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.532 `(config.encoder_attention_heads,)`.533 output_attentions (`bool`, *optional*):534 Whether or not to return the attentions tensors of all attention layers. See `attentions` under535 returned tensors for more detail.536 """537 residual = hidden_states538 539 hidden_states = self.layer_norm1(hidden_states)540 hidden_states, attn_weights = self.self_attn(541 hidden_states=hidden_states,542 attention_mask=attention_mask,543 causal_attention_mask=causal_attention_mask,544 output_attentions=output_attentions,545 )546 hidden_states = residual + hidden_states547 548 residual = hidden_states549 hidden_states = self.layer_norm2(hidden_states)550 hidden_states = self.mlp(hidden_states)551 hidden_states = residual + hidden_states552 553 outputs = (hidden_states,)554 555 if output_attentions:556 outputs += (attn_weights,)557 558 return outputs559 560 561@auto_docstring562# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTPreTrainedModel with OwlViT->Owlv2,owlvit->owlv2563class Owlv2PreTrainedModel(PreTrainedModel):564 config: Owlv2Config565 base_model_prefix = "owlv2"566 supports_gradient_checkpointing = True567 _no_split_modules = ["Owlv2EncoderLayer"]568 569 def _init_weights(self, module: nn.Module):570 """Initialize the weights"""571 factor = self.config.initializer_factor572 if isinstance(module, Owlv2TextEmbeddings):573 module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)574 module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)575 elif isinstance(module, Owlv2VisionEmbeddings):576 nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)577 nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)578 nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)579 elif isinstance(module, Owlv2Attention):580 in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor581 out_proj_std = (module.embed_dim**-0.5) * factor582 nn.init.normal_(module.q_proj.weight, std=in_proj_std)583 nn.init.normal_(module.k_proj.weight, std=in_proj_std)584 nn.init.normal_(module.v_proj.weight, std=in_proj_std)585 nn.init.normal_(module.out_proj.weight, std=out_proj_std)586 elif isinstance(module, Owlv2MLP):587 in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor588 fc_std = (2 * module.config.hidden_size) ** -0.5 * factor589 nn.init.normal_(module.fc1.weight, std=fc_std)590 nn.init.normal_(module.fc2.weight, std=in_proj_std)591 elif isinstance(module, Owlv2Model):592 nn.init.normal_(593 module.text_projection.weight,594 std=module.text_embed_dim**-0.5 * factor,595 )596 nn.init.normal_(597 module.visual_projection.weight,598 std=module.vision_embed_dim**-0.5 * factor,599 )600 module.logit_scale.data.fill_(self.config.logit_scale_init_value)601 if isinstance(module, nn.LayerNorm):602 module.bias.data.zero_()603 module.weight.data.fill_(1.0)604 if isinstance(module, nn.Linear):605 module.weight.data.normal_(mean=0.0, std=factor)606 if module.bias is not None:607 module.bias.data.zero_()608 609 610# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTEncoder with OwlViT->Owlv2611class Owlv2Encoder(nn.Module):612 """613 Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a614 [`Owlv2EncoderLayer`].615 616 Args:617 config: Owlv2Config618 """619 620 def __init__(self, config: Owlv2Config):621 super().__init__()622 self.layers = nn.ModuleList([Owlv2EncoderLayer(config) for _ in range(config.num_hidden_layers)])623 self.gradient_checkpointing = False624 625 def forward(626 self,627 inputs_embeds,628 attention_mask: Optional[torch.Tensor] = None,629 causal_attention_mask: Optional[torch.Tensor] = None,630 output_attentions: Optional[bool] = None,631 output_hidden_states: Optional[bool] = None,632 return_dict: Optional[bool] = None,633 ) -> Union[tuple, BaseModelOutput]:634 r"""635 Args:636 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`).637 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):638 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:639 - 1 for tokens that are **not masked**,640 - 0 for tokens that are **masked**.641 [What are attention masks?](../glossary#attention-mask)642 causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):643 Causal mask for the text model. Mask values selected in `[0, 1]`:644 - 1 for tokens that are **not masked**,645 - 0 for tokens that are **masked**.646 [What are attention masks?](../glossary#attention-mask)647 output_attentions (`bool`, *optional*):648 Whether or not to return the attentions tensors of all attention layers. See `attentions` under649 returned tensors for more detail.650 output_hidden_states (`bool`, *optional*):651 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors652 for more detail.653 return_dict (`bool`, *optional*):654 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.655 """656 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions657 output_hidden_states = (658 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states659 )660 return_dict = return_dict if return_dict is not None else self.config.use_return_dict661 662 encoder_states = () if output_hidden_states else None663 all_attentions = () if output_attentions else None664 665 hidden_states = inputs_embeds666 for encoder_layer in self.layers:667 if output_hidden_states:668 encoder_states = encoder_states + (hidden_states,)669 layer_outputs = encoder_layer(670 hidden_states,671 attention_mask,672 causal_attention_mask,673 output_attentions=output_attentions,674 )675 676 hidden_states = layer_outputs[0]677 678 if output_attentions:679 all_attentions = all_attentions + (layer_outputs[1],)680 681 if output_hidden_states:682 encoder_states = encoder_states + (hidden_states,)683 684 if not return_dict:685 return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)686 return BaseModelOutput(687 last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions688 )689 690 691# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextTransformer with OWLVIT->OWLV2,OwlViT->Owlv2692class Owlv2TextTransformer(nn.Module):693 def __init__(self, config: Owlv2TextConfig):694 super().__init__()695 self.config = config696 embed_dim = config.hidden_size697 self.embeddings = Owlv2TextEmbeddings(config)698 self.encoder = Owlv2Encoder(config)699 self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)700 701 @auto_docstring702 def forward(703 self,704 input_ids: torch.Tensor,705 attention_mask: Optional[torch.Tensor] = None,706 position_ids: Optional[torch.Tensor] = None,707 output_attentions: Optional[bool] = None,708 output_hidden_states: Optional[bool] = None,709 return_dict: Optional[bool] = None,710 ) -> Union[tuple, BaseModelOutputWithPooling]:711 r"""712 input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`):713 Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See714 [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input715 IDs?](../glossary#input-ids)716 """717 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions718 output_hidden_states = (719 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states720 )721 return_dict = return_dict if return_dict is not None else self.config.use_return_dict722 723 input_shape = input_ids.size()724 input_ids = input_ids.view(-1, input_shape[-1])725 hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)726 727 # num_samples, seq_len = input_shape where num_samples = batch_size * num_max_text_queries728 # OWLV2's text model uses causal mask, prepare it here.729 # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324730 causal_attention_mask = _create_4d_causal_attention_mask(731 input_shape, hidden_states.dtype, device=hidden_states.device732 )733 # expand attention_mask734 if attention_mask is not None:735 # [num_samples, seq_len] -> [num_samples, 1, tgt_seq_len, src_seq_len]736 attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)737 738 encoder_outputs = self.encoder(739 inputs_embeds=hidden_states,740 attention_mask=attention_mask,741 causal_attention_mask=causal_attention_mask,742 output_attentions=output_attentions,743 output_hidden_states=output_hidden_states,744 return_dict=return_dict,745 )746 747 last_hidden_state = encoder_outputs[0]748 last_hidden_state = self.final_layer_norm(last_hidden_state)749 750 # take features from the end of tokens embedding (end of token is the highest number in each sequence)751 # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14752 pooled_output = last_hidden_state[753 torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),754 input_ids.to(torch.int).argmax(dim=-1).to(last_hidden_state.device),755 ]756 757 if not return_dict:758 return (last_hidden_state, pooled_output) + encoder_outputs[1:]759 760 return BaseModelOutputWithPooling(761 last_hidden_state=last_hidden_state,762 pooler_output=pooled_output,763 hidden_states=encoder_outputs.hidden_states,764 attentions=encoder_outputs.attentions,765 )766 767 768# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTTextModel with google/owlvit-base-patch32->google/owlv2-base-patch16, OWLVIT->OWLV2,OwlViT->Owlv2769class Owlv2TextModel(Owlv2PreTrainedModel):770 config: Owlv2TextConfig771 772 def __init__(self, config: Owlv2TextConfig):773 super().__init__(config)774 self.text_model = Owlv2TextTransformer(config)775 # Initialize weights and apply final processing776 self.post_init()777 778 def get_input_embeddings(self) -> nn.Module:779 return self.text_model.embeddings.token_embedding780 781 def set_input_embeddings(self, value):782 self.text_model.embeddings.token_embedding = value783 784 @auto_docstring785 def forward(786 self,787 input_ids: torch.Tensor,788 attention_mask: Optional[torch.Tensor] = None,789 output_attentions: Optional[bool] = None,790 output_hidden_states: Optional[bool] = None,791 return_dict: Optional[bool] = None,792 ) -> Union[tuple, BaseModelOutputWithPooling]:793 r"""794 input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`):795 Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See796 [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input797 IDs?](../glossary#input-ids)798 799 Examples:800 ```python801 >>> from transformers import AutoProcessor, Owlv2TextModel802 803 >>> model = Owlv2TextModel.from_pretrained("google/owlv2-base-patch16")804 >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16")805 >>> inputs = processor(806 ... text=[["a photo of a cat", "a photo of a dog"], ["photo of a astranaut"]], return_tensors="pt"807 ... )808 >>> outputs = model(**inputs)809 >>> last_hidden_state = outputs.last_hidden_state810 >>> pooled_output = outputs.pooler_output # pooled (EOS token) states811 ```"""812 813 # Get embeddings for all text queries in all batch samples814 return self.text_model(815 input_ids=input_ids,816 attention_mask=attention_mask,817 output_attentions=output_attentions,818 output_hidden_states=output_hidden_states,819 return_dict=return_dict,820 )821 822 823# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionTransformer with OWLVIT->OWLV2,OwlViT->Owlv2824class Owlv2VisionTransformer(nn.Module):825 def __init__(self, config: Owlv2VisionConfig):826 super().__init__()827 self.config = config828 829 self.embeddings = Owlv2VisionEmbeddings(config)830 self.pre_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)831 self.encoder = Owlv2Encoder(config)832 self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)833 834 @auto_docstring835 def forward(836 self,837 pixel_values: torch.FloatTensor,838 output_attentions: Optional[bool] = None,839 output_hidden_states: Optional[bool] = None,840 interpolate_pos_encoding: Optional[bool] = False,841 return_dict: Optional[bool] = None,842 ) -> Union[tuple, BaseModelOutputWithPooling]:843 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions844 output_hidden_states = (845 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states846 )847 return_dict = return_dict if return_dict is not None else self.config.use_return_dict848 849 # Cast the input to the expected `dtype`850 expected_input_dtype = self.embeddings.patch_embedding.weight.dtype851 pixel_values = pixel_values.to(expected_input_dtype)852 853 hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)854 hidden_states = self.pre_layernorm(hidden_states)855 856 encoder_outputs = self.encoder(857 inputs_embeds=hidden_states,858 output_attentions=output_attentions,859 output_hidden_states=output_hidden_states,860 return_dict=return_dict,861 )862 863 last_hidden_state = encoder_outputs[0]864 pooled_output = last_hidden_state[:, 0, :]865 866 pooled_output = self.post_layernorm(pooled_output)867 868 if not return_dict:869 return (last_hidden_state, pooled_output) + encoder_outputs[1:]870 871 return BaseModelOutputWithPooling(872 last_hidden_state=last_hidden_state,873 pooler_output=pooled_output,874 hidden_states=encoder_outputs.hidden_states,875 attentions=encoder_outputs.attentions,876 )877 878 879# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTVisionModel with OWLVIT->OWLV2,OwlViT->Owlv2,google/owlvit-base-patch32->google/owlv2-base-patch16880class Owlv2VisionModel(Owlv2PreTrainedModel):881 config: Owlv2VisionConfig882 main_input_name = "pixel_values"883 884 def __init__(self, config: Owlv2VisionConfig):885 super().__init__(config)886 self.vision_model = Owlv2VisionTransformer(config)887 # Initialize weights and apply final processing888 self.post_init()889 890 def get_input_embeddings(self) -> nn.Module:891 return self.vision_model.embeddings.patch_embedding892 893 @auto_docstring894 def forward(895 self,896 pixel_values: Optional[torch.FloatTensor] = None,897 output_attentions: Optional[bool] = None,898 output_hidden_states: Optional[bool] = None,899 interpolate_pos_encoding: bool = False,900 return_dict: Optional[bool] = None,901 ) -> Union[tuple, BaseModelOutputWithPooling]:902 r"""903 Examples:904 ```python905 >>> from PIL import Image906 >>> import requests907 >>> from transformers import AutoProcessor, Owlv2VisionModel908 909 >>> model = Owlv2VisionModel.from_pretrained("google/owlv2-base-patch16")910 >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16")911 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"912 >>> image = Image.open(requests.get(url, stream=True).raw)913 914 >>> inputs = processor(images=image, return_tensors="pt")915 916 >>> outputs = model(**inputs)917 >>> last_hidden_state = outputs.last_hidden_state918 >>> pooled_output = outputs.pooler_output # pooled CLS states919 ```"""920 return self.vision_model(921 pixel_values=pixel_values,922 output_attentions=output_attentions,923 output_hidden_states=output_hidden_states,924 interpolate_pos_encoding=interpolate_pos_encoding,925 return_dict=return_dict,926 )927 928 929@auto_docstring930# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTModel with google/owlvit-base-patch32->google/owlv2-base-patch16-ensemble, OWLVIT->OWLV2,OwlViT->Owlv2,owlvit->owlv2,OWL-ViT->OWLv2931class Owlv2Model(Owlv2PreTrainedModel):932 config: Owlv2Config933 934 def __init__(self, config: Owlv2Config):935 super().__init__(config)936 937 if not isinstance(config.text_config, Owlv2TextConfig):938 raise TypeError(939 "config.text_config is expected to be of type Owlv2TextConfig but is of type"940 f" {type(config.text_config)}."941 )942 943 if not isinstance(config.vision_config, Owlv2VisionConfig):944 raise TypeError(945 "config.vision_config is expected to be of type Owlv2VisionConfig but is of type"946 f" {type(config.vision_config)}."947 )948 949 text_config = config.text_config950 vision_config = config.vision_config951 952 self.projection_dim = config.projection_dim953 self.text_embed_dim = text_config.hidden_size954 self.vision_embed_dim = vision_config.hidden_size955 956 self.text_model = Owlv2TextTransformer(text_config)957 self.vision_model = Owlv2VisionTransformer(vision_config)958 959 self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)960 self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)961 self.logit_scale = nn.Parameter(torch.tensor(config.logit_scale_init_value))962 963 # Initialize weights and apply final processing964 self.post_init()965 966 @filter_out_non_signature_kwargs()967 @auto_docstring968 def get_text_features(969 self,970 input_ids: torch.Tensor,971 attention_mask: Optional[torch.Tensor] = None,972 ) -> torch.FloatTensor:973 r"""974 input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`):975 Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See976 [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input977 IDs?](../glossary#input-ids)978 979 Returns:980 text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by981 applying the projection layer to the pooled output of [`Owlv2TextModel`].982 983 Examples:984 ```python985 >>> import torch986 >>> from transformers import AutoProcessor, Owlv2Model987 988 >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")989 >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")990 >>> inputs = processor(991 ... text=[["a photo of a cat", "a photo of a dog"], ["photo of a astranaut"]], return_tensors="pt"992 ... )993 >>> with torch.inference_mode():994 ... text_features = model.get_text_features(**inputs)995 ```"""996 # Get embeddings for all text queries in all batch samples997 text_outputs: BaseModelOutputWithPooling = self.text_model(input_ids=input_ids, attention_mask=attention_mask)998 text_features = self.text_projection(text_outputs.pooler_output)999 1000 return text_features1001 1002 @filter_out_non_signature_kwargs()1003 @auto_docstring1004 def get_image_features(1005 self,1006 pixel_values: torch.Tensor,1007 interpolate_pos_encoding: bool = False,1008 ) -> torch.FloatTensor:1009 r"""1010 Returns:1011 image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by1012 applying the projection layer to the pooled output of [`Owlv2VisionModel`].1013 1014 Examples:1015 ```python1016 >>> import torch1017 >>> from transformers.image_utils import load_image1018 >>> from transformers import AutoProcessor, Owlv2Model1019 1020 >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")1021 >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")1022 1023 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"1024 >>> image = load_image(url)1025 1026 >>> inputs = processor(images=image, return_tensors="pt")1027 >>> with torch.inference_mode():1028 ... image_features = model.get_image_features(**inputs)1029 ```"""1030 vision_outputs: BaseModelOutputWithPooling = self.vision_model(1031 pixel_values=pixel_values,1032 interpolate_pos_encoding=interpolate_pos_encoding,1033 )1034 image_features = self.visual_projection(vision_outputs.pooler_output)1035 1036 return image_features1037 1038 @auto_docstring1039 def forward(1040 self,1041 input_ids: Optional[torch.LongTensor] = None,1042 pixel_values: Optional[torch.FloatTensor] = None,1043 attention_mask: Optional[torch.Tensor] = None,1044 return_loss: Optional[bool] = None,1045 output_attentions: Optional[bool] = None,1046 output_hidden_states: Optional[bool] = None,1047 interpolate_pos_encoding: bool = False,1048 return_base_image_embeds: Optional[bool] = None,1049 return_dict: Optional[bool] = None,1050 ) -> Union[tuple, Owlv2Output]:1051 r"""1052 return_loss (`bool`, *optional*):1053 Whether or not to return the contrastive loss.1054 return_base_image_embeds (`bool`, *optional*):1055 Whether or not to return the base image embeddings.1056 1057 Examples:1058 ```python1059 >>> from PIL import Image1060 >>> import requests1061 >>> from transformers import AutoProcessor, Owlv2Model1062 1063 >>> model = Owlv2Model.from_pretrained("google/owlv2-base-patch16-ensemble")1064 >>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")1065 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"1066 >>> image = Image.open(requests.get(url, stream=True).raw)1067 >>> inputs = processor(text=[["a photo of a cat", "a photo of a dog"]], images=image, return_tensors="pt")1068 >>> outputs = model(**inputs)1069 >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score1070 >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities1071 ```"""1072 # Use OWLv2 model's config for some fields (if specified) instead of those of vision & text components.1073 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1074 output_hidden_states = (1075 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1076 )1077 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1078 1079 vision_outputs = self.vision_model(1080 pixel_values=pixel_values,1081 output_attentions=output_attentions,1082 output_hidden_states=output_hidden_states,1083 interpolate_pos_encoding=interpolate_pos_encoding,1084 return_dict=return_dict,1085 )1086 1087 # Get embeddings for all text queries in all batch samples1088 text_outputs = self.text_model(1089 input_ids=input_ids,1090 attention_mask=attention_mask,1091 output_attentions=output_attentions,1092 output_hidden_states=output_hidden_states,1093 return_dict=return_dict,1094 )1095 1096 text_embeds = text_outputs[1]1097 text_embeds = self.text_projection(text_embeds)1098 image_embeds = vision_outputs[1]1099 image_embeds = self.visual_projection(image_embeds)1100 1101 # normalized features1102 image_embeds = image_embeds / torch.linalg.norm(image_embeds, ord=2, dim=-1, keepdim=True)1103 text_embeds_norm = text_embeds / torch.linalg.norm(text_embeds, ord=2, dim=-1, keepdim=True)1104 1105 # cosine similarity as logits and set it on the correct device1106 logit_scale = self.logit_scale.exp().to(image_embeds.device)1107 1108 logits_per_text = torch.matmul(text_embeds_norm, image_embeds.t()) * logit_scale1109 logits_per_image = logits_per_text.t()1110 1111 loss = None1112 if return_loss:1113 loss = owlv2_loss(logits_per_text)1114 1115 text_embeds = text_embeds_norm1116 1117 if not return_dict:1118 output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)1119 return ((loss,) + output) if loss is not None else output1120 1121 return Owlv2Output(1122 loss=loss,1123 logits_per_image=logits_per_image,1124 logits_per_text=logits_per_text,1125 text_embeds=text_embeds,1126 image_embeds=image_embeds,1127 text_model_output=text_outputs,1128 vision_model_output=vision_outputs,1129 )1130 1131 1132# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTBoxPredictionHead with OwlViT->Owlv21133class Owlv2BoxPredictionHead(nn.Module):1134 def __init__(self, config: Owlv2Config, out_dim: int = 4):1135 super().__init__()1136 1137 width = config.vision_config.hidden_size1138 self.dense0 = nn.Linear(width, width)1139 self.dense1 = nn.Linear(width, width)1140 self.gelu = nn.GELU()1141 self.dense2 = nn.Linear(width, out_dim)1142 1143 def forward(self, image_features: torch.Tensor) -> torch.FloatTensor:1144 output = self.dense0(image_features)1145 output = self.gelu(output)1146 output = self.dense1(output)1147 output = self.gelu(output)1148 output = self.dense2(output)1149 return output1150 1151 1152# Copied from transformers.models.owlvit.modeling_owlvit.OwlViTClassPredictionHead with OwlViT->Owlv21153class Owlv2ClassPredictionHead(nn.Module):1154 def __init__(self, config: Owlv2Config):1155 super().__init__()1156 1157 out_dim = config.text_config.hidden_size1158 self.query_dim = config.vision_config.hidden_size1159 1160 self.dense0 = nn.Linear(self.query_dim, out_dim)1161 self.logit_shift = nn.Linear(self.query_dim, 1)1162 self.logit_scale = nn.Linear(self.query_dim, 1)1163 self.elu = nn.ELU()1164 1165 def forward(1166 self,1167 image_embeds: torch.FloatTensor,1168 query_embeds: Optional[torch.FloatTensor],1169 query_mask: Optional[torch.Tensor],1170 ) -> tuple[torch.FloatTensor]:1171 image_class_embeds = self.dense0(image_embeds)1172 if query_embeds is None:1173 device = image_class_embeds.device1174 batch_size, num_patches = image_class_embeds.shape[:2]1175 pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)1176 return (pred_logits, image_class_embeds)1177 1178 # Normalize image and text features1179 image_class_embeds = image_class_embeds / (torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6)1180 query_embeds = query_embeds / (torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6)1181 1182 # Get class predictions1183 pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)1184 1185 # Apply a learnable shift and scale to logits1186 logit_shift = self.logit_shift(image_embeds)1187 logit_scale = self.logit_scale(image_embeds)1188 logit_scale = self.elu(logit_scale) + 11189 pred_logits = (pred_logits + logit_shift) * logit_scale1190 1191 if query_mask is not None:1192 if query_mask.ndim > 1:1193 query_mask = torch.unsqueeze(query_mask, dim=-2)1194 1195 pred_logits = torch.where(query_mask == 0, torch.finfo(pred_logits.dtype).min, pred_logits)1196 pred_logits = pred_logits.to(torch.float32)1197 1198 return (pred_logits, image_class_embeds)1199 1200