
import torch
from torch import nn



class 		UpperCamelCase__  (       nn.Module    ):
   def __init__(self :					Optional[int]					,				snake_case_ :					int					,				snake_case_ :					Optional[int]					,				snake_case_ :					Dict					,				snake_case_ :					Dict					,				snake_case_ :					List[Any]=1					,				snake_case_ :					Any=False   ):
    super().__init__()

    __a :  int						=  n_token
    __a :  Optional[int]						=  d_embed
    __a :  Dict						=  d_proj

    __a :  Optional[Any]						=  cutoffs + [n_token]
    __a :  Tuple						=  [0] + self.cutoffs
    __a :  List[Any]						=  div_val

    __a :  Optional[Any]						=  self.cutoffs[0]
    __a :  List[str]						=  len(self.cutoffs   ) - 1
    __a :  Optional[int]						=  self.shortlist_size + self.n_clusters

    if self.n_clusters > 0:
     __a :  Any						=  nn.Parameter(torch.zeros(self.n_clusters					,				self.d_embed   )   )
     __a :  Optional[Any]						=  nn.Parameter(torch.zeros(self.n_clusters   )   )

    __a :  Optional[Any]						=  nn.ModuleList()
    __a :  Tuple						=  nn.ParameterList()

    if div_val == 1:
     for i in range(len(self.cutoffs   )   ):
      if d_proj != d_embed:
       self.out_projs.append(nn.Parameter(torch.FloatTensor(snake_case_					,				snake_case_   )   )   )
      else:
       self.out_projs.append(snake_case_   )

     self.out_layers.append(nn.Linear(snake_case_					,				snake_case_   )   )
    else:
     for i in range(len(self.cutoffs   )   ):
      __a				,       __a :  str						=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __a :  Union[str, Any]						=  d_embed // (div_val**i)

      self.out_projs.append(nn.Parameter(torch.FloatTensor(snake_case_					,				snake_case_   )   )   )

      self.out_layers.append(nn.Linear(snake_case_					,				r_idx - l_idx   )   )

    __a :  Optional[int]						=  keep_order
   def       lowerCAmelCase			(self :					str					,				snake_case_ :					Any					,				snake_case_ :					Union[str, Any]					,				snake_case_ :					Any					,				snake_case_ :					str   ):
    if proj is None:
     __a :  Optional[Any]						=  nn.functional.linear(snake_case_					,				snake_case_					,				bias=snake_case_   )
    else:
     # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
     __a :  Any						=  nn.functional.linear(snake_case_					,				proj.t().contiguous()   )
     __a :  List[str]						=  nn.functional.linear(snake_case_					,				snake_case_					,				bias=snake_case_   )
     # else:
     #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
     #     if bias is not None:
     #         logit = logit + bias

    return logit
   def       lowerCAmelCase			(self :					Optional[int]					,				snake_case_ :					Tuple					,				snake_case_ :					Union[str, Any]=None					,				snake_case_ :					Any=False   ):

    if labels is not None:
     # Shift so that tokens < n predict n
     __a :  Optional[int]						=  hidden[..., :-1, :].contiguous()
     __a :  Union[str, Any]						=  labels[..., 1:].contiguous()
     __a :  Optional[Any]						=  hidden.view(-1					,				hidden.size(-1   )   )
     __a :  Union[str, Any]						=  labels.view(-1   )
     if hidden.size(0   ) != labels.size(0   ):
      raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''   )
    else:
     __a :  Optional[int]						=  hidden.view(-1					,				hidden.size(-1   )   )

    if self.n_clusters == 0:
     __a :  List[Any]						=  self._compute_logit(snake_case_					,				self.out_layers[0].weight					,				self.out_layers[0].bias					,				self.out_projs[0]   )
     if labels is not None:
      __a :  int						=  labels != -1_0_0
      __a :  int						=  torch.zeros_like(snake_case_					,				dtype=hidden.dtype					,				device=hidden.device   )
      __a :  Optional[Any]						=  (
          -nn.functional.log_softmax(snake_case_					,				dim=-1   )[mask].gather(1					,				labels[mask].unsqueeze(1   )   ).squeeze(1   )
      )
     else:
      __a :  Optional[int]						=  nn.functional.log_softmax(snake_case_					,				dim=-1   )
    else:
     # construct weights and biases
     __a				,       __a :  Tuple						=  [], []
     for i in range(len(self.cutoffs   )   ):
      if self.div_val == 1:
       __a				,       __a :  int						=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
       __a :  List[str]						=  self.out_layers[0].weight[l_idx:r_idx]
       __a :  Union[str, Any]						=  self.out_layers[0].bias[l_idx:r_idx]
      else:
       __a :  Union[str, Any]						=  self.out_layers[i].weight
       __a :  Tuple						=  self.out_layers[i].bias

      if i == 0:
       __a :  int						=  torch.cat([weight_i, self.cluster_weight]					,				dim=0   )
       __a :  int						=  torch.cat([bias_i, self.cluster_bias]					,				dim=0   )

      weights.append(snake_case_   )
      biases.append(snake_case_   )

     __a				,       __a				,       __a :  List[str]						=  weights[0], biases[0], self.out_projs[0]

     __a :  Optional[Any]						=  self._compute_logit(snake_case_					,				snake_case_					,				snake_case_					,				snake_case_   )
     __a :  Union[str, Any]						=  nn.functional.log_softmax(snake_case_					,				dim=1   )

     if labels is None:
      __a :  str						=  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
     else:
      __a :  str						=  torch.zeros_like(snake_case_					,				dtype=hidden.dtype					,				device=hidden.device   )

     __a :  Union[str, Any]						=  0
     __a :  Optional[int]						=  [0] + self.cutoffs
     for i in range(len(snake_case_   ) - 1   ):
      __a				,       __a :  Tuple						=  cutoff_values[i], cutoff_values[i + 1]

      if labels is not None:
       __a :  List[str]						=  (labels >= l_idx) & (labels < r_idx)
       __a :  List[Any]						=  mask_i.nonzero().squeeze()

       if indices_i.numel() == 0:
        continue

       __a :  Optional[Any]						=  labels.index_select(0					,				snake_case_   ) - l_idx
       __a :  Dict						=  head_logprob.index_select(0					,				snake_case_   )
       __a :  List[Any]						=  hidden.index_select(0					,				snake_case_   )
      else:
       __a :  Union[str, Any]						=  hidden

      if i == 0:
       if labels is not None:
        __a :  List[str]						=  head_logprob_i.gather(1					,				target_i[:, None]   ).squeeze(1   )
       else:
        __a :  Optional[Any]						=  head_logprob[:, : self.cutoffs[0]]
      else:
       __a				,       __a				,       __a :  List[Any]						=  weights[i], biases[i], self.out_projs[i]

       __a :  Dict						=  self._compute_logit(snake_case_					,				snake_case_					,				snake_case_					,				snake_case_   )
       __a :  Optional[int]						=  nn.functional.log_softmax(snake_case_					,				dim=1   )
       __a :  Optional[Any]						=  self.cutoffs[0] + i - 1  # No probability for the head cluster
       if labels is not None:
        __a :  Dict						=  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
            1					,				target_i[:, None]   ).squeeze(1   )
       else:
        __a :  Optional[int]						=  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
        __a :  str						=  logprob_i

      if labels is not None:
       if (hasattr(self					,				'''keep_order'''   ) and self.keep_order) or keep_order:
        out.index_copy_(0					,				snake_case_					,				-logprob_i   )
       else:
        out[offset : offset + logprob_i.size(0   )].copy_(-logprob_i   )
       offset += logprob_i.size(0   )

    return out







   def       lowerCAmelCase			(self :					List[str]					,				snake_case_ :					Dict   ):
    if self.n_clusters == 0:
     __a :  Dict						=  self._compute_logit(snake_case_					,				self.out_layers[0].weight					,				self.out_layers[0].bias					,				self.out_projs[0]   )
     return nn.functional.log_softmax(snake_case_					,				dim=-1   )
    else:
     # construct weights and biases
     __a				,       __a :  Any						=  [], []
     for i in range(len(self.cutoffs   )   ):
      if self.div_val == 1:
       __a				,       __a :  Optional[Any]						=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
       __a :  Tuple						=  self.out_layers[0].weight[l_idx:r_idx]
       __a :  Dict						=  self.out_layers[0].bias[l_idx:r_idx]
      else:
       __a :  Any						=  self.out_layers[i].weight
       __a :  List[Any]						=  self.out_layers[i].bias

      if i == 0:
       __a :  Dict						=  torch.cat([weight_i, self.cluster_weight]					,				dim=0   )
       __a :  Dict						=  torch.cat([bias_i, self.cluster_bias]					,				dim=0   )

      weights.append(snake_case_   )
      biases.append(snake_case_   )

     __a				,       __a				,       __a :  Dict						=  weights[0], biases[0], self.out_projs[0]
     __a :  List[str]						=  self._compute_logit(snake_case_					,				snake_case_					,				snake_case_					,				snake_case_   )

     __a :  Tuple						=  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
     __a :  Optional[Any]						=  nn.functional.log_softmax(snake_case_					,				dim=1   )

     __a :  str						=  [0] + self.cutoffs
     for i in range(len(snake_case_   ) - 1   ):
      __a				,       __a :  Any						=  cutoff_values[i], cutoff_values[i + 1]

      if i == 0:
       __a :  str						=  head_logprob[:, : self.cutoffs[0]]
      else:
       __a				,       __a				,       __a :  Dict						=  weights[i], biases[i], self.out_projs[i]

       __a :  Optional[int]						=  self._compute_logit(snake_case_					,				snake_case_					,				snake_case_					,				snake_case_   )
       __a :  Optional[int]						=  nn.functional.log_softmax(snake_case_					,				dim=1   )

       __a :  Optional[Any]						=  head_logprob[:, -i] + tail_logprob_i
       __a :  Union[str, Any]						=  logprob_i

     return out


import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel




def    __UpperCamelCase     (				lowerCAmelCase__							:				Any  ):
 # vision encoder
 if "img_encoder.pos_embed" in name:
  __a :  Any						=  name.replace('''img_encoder.pos_embed'''  ,   '''vision_model.embeddings.position_embeddings'''  )
 if "img_encoder.patch_embed.proj" in name:
  __a :  str						=  name.replace('''img_encoder.patch_embed.proj'''  ,   '''vision_model.embeddings.patch_embeddings.projection'''  )
 if "img_encoder.patch_embed.norm" in name:
  __a :  int						=  name.replace('''img_encoder.patch_embed.norm'''  ,   '''vision_model.embeddings.layernorm'''  )
 if "img_encoder.layers" in name:
  __a :  Union[str, Any]						=  name.replace('''img_encoder.layers'''  ,   '''vision_model.encoder.stages'''  )
 if "blocks" in name and "res" not in name:
  __a :  List[Any]						=  name.replace('''blocks'''  ,   '''layers'''  )
 if "attn" in name and "pre_assign" not in name:
  __a :  Tuple						=  name.replace('''attn'''  ,   '''self_attn'''  )
 if "proj" in name and "self_attn" in name and "text" not in name:
  __a :  List[Any]						=  name.replace('''proj'''  ,   '''out_proj'''  )
 if "pre_assign_attn.attn.proj" in name:
  __a :  Any						=  name.replace('''pre_assign_attn.attn.proj'''  ,   '''pre_assign_attn.attn.out_proj'''  )
 if "norm1" in name:
  __a :  Union[str, Any]						=  name.replace('''norm1'''  ,   '''layer_norm1'''  )
 if "norm2" in name and "pre_assign" not in name:
  __a :  Optional[int]						=  name.replace('''norm2'''  ,   '''layer_norm2'''  )
 if "img_encoder.norm" in name:
  __a :  Union[str, Any]						=  name.replace('''img_encoder.norm'''  ,   '''vision_model.layernorm'''  )
 # text encoder
 if "text_encoder.token_embedding" in name:
  __a :  List[Any]						=  name.replace('''text_encoder.token_embedding'''  ,   '''text_model.embeddings.token_embedding'''  )
 if "text_encoder.positional_embedding" in name:
  __a :  Any						=  name.replace('''text_encoder.positional_embedding'''  ,   '''text_model.embeddings.position_embedding.weight'''  )
 if "text_encoder.transformer.resblocks." in name:
  __a :  Any						=  name.replace('''text_encoder.transformer.resblocks.'''  ,   '''text_model.encoder.layers.'''  )
 if "ln_1" in name:
  __a :  str						=  name.replace('''ln_1'''  ,   '''layer_norm1'''  )
 if "ln_2" in name:
  __a :  Union[str, Any]						=  name.replace('''ln_2'''  ,   '''layer_norm2'''  )
 if "c_fc" in name:
  __a :  Union[str, Any]						=  name.replace('''c_fc'''  ,   '''fc1'''  )
 if "c_proj" in name:
  __a :  Union[str, Any]						=  name.replace('''c_proj'''  ,   '''fc2'''  )
 if "text_encoder" in name:
  __a :  Optional[int]						=  name.replace('''text_encoder'''  ,   '''text_model'''  )
 if "ln_final" in name:
  __a :  str						=  name.replace('''ln_final'''  ,   '''final_layer_norm'''  )
 # projection layers
 if "img_projector.linear_hidden." in name:
  __a :  List[str]						=  name.replace('''img_projector.linear_hidden.'''  ,   '''visual_projection.'''  )
 if "img_projector.linear_out." in name:
  __a :  str						=  name.replace('''img_projector.linear_out.'''  ,   '''visual_projection.3.'''  )
 if "text_projector.linear_hidden" in name:
  __a :  int						=  name.replace('''text_projector.linear_hidden'''  ,   '''text_projection'''  )
 if "text_projector.linear_out" in name:
  __a :  List[str]						=  name.replace('''text_projector.linear_out'''  ,   '''text_projection.3'''  )

 return name




def    __UpperCamelCase     (				lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				Tuple  ):
 for key in orig_state_dict.copy().keys():
  __a :  List[Any]						=  orig_state_dict.pop(lowerCAmelCase__  )

  if "qkv" in key:
   # weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
   # we need to split them up into separate matrices/vectors
   __a :  Tuple						=  key.split('''.'''  )
   __a				,       __a :  List[Any]						=  int(key_split[2]  ), int(key_split[4]  )
   __a :  List[Any]						=  config.vision_config.hidden_size
   if "weight" in key:
    __a :  int						=  val[:dim, :]
    __a :  List[str]						=  val[dim : dim * 2, :]
    __a :  List[Any]						=  val[-dim:, :]
   else:
    __a :  List[str]						=  val[:dim]
    __a :  int						=  val[dim : dim * 2]
    __a :  Any						=  val[-dim:]
  elif "in_proj" in key:
   # weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
   # we need to split them up into separate matrices/vectors
   __a :  int						=  key.split('''.'''  )
   __a :  str						=  int(key_split[3]  )
   __a :  List[Any]						=  config.text_config.hidden_size
   if "weight" in key:
    __a :  List[str]						=  val[:dim, :]
    __a :  Any						=  val[
        dim : dim * 2, :
    ]
    __a :  Dict						=  val[-dim:, :]
   else:
    __a :  List[str]						=  val[:dim]
    __a :  Any						=  val[dim : dim * 2]
    __a :  Any						=  val[-dim:]
  else:
   __a :  Union[str, Any]						=  rename_key(lowerCAmelCase__  )
   # squeeze if necessary
   if (
       "text_projection.0" in new_name
       or "text_projection.3" in new_name
       or "visual_projection.0" in new_name
       or "visual_projection.3" in new_name
   ):
    __a :  List[Any]						=  val.squeeze_()
   else:
    __a :  Dict						=  val

 return orig_state_dict




def    __UpperCamelCase     (				):
 __a :  Optional[int]						=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 __a :  str						=  Image.open(requests.get(lowerCAmelCase__  ,   stream=lowerCAmelCase__  ).raw  )
 return im




@torch.no_grad()
def    __UpperCamelCase     (				lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				Optional[Any]="groupvit-gcc-yfcc"  ,   lowerCAmelCase__							:				int=False  ):
 __a :  Union[str, Any]						=  GroupViTConfig()
 __a :  int						=  GroupViTModel(lowerCAmelCase__  ).eval()

 __a :  Any						=  torch.load(lowerCAmelCase__  ,   map_location='''cpu'''  )['''model''']
 __a :  Optional[Any]						=  convert_state_dict(lowerCAmelCase__  ,   lowerCAmelCase__  )
 __a				,       __a :  Dict						=  model.load_state_dict(lowerCAmelCase__  ,   strict=lowerCAmelCase__  )
 assert missing_keys == ["text_model.embeddings.position_ids"]
 assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(lowerCAmelCase__  ) == 0)

 # verify result
 __a :  Any						=  CLIPProcessor.from_pretrained('''openai/clip-vit-base-patch32'''  )
 __a :  Optional[Any]						=  prepare_img()
 __a :  Optional[int]						=  processor(text=['''a photo of a cat''', '''a photo of a dog''']  ,   images=lowerCAmelCase__  ,   padding=lowerCAmelCase__  ,   return_tensors='''pt'''  )

 with torch.no_grad():
  __a :  Tuple						=  model(**lowerCAmelCase__  )

 if model_name == "groupvit-gcc-yfcc":
  __a :  List[str]						=  torch.tensor([[13.35_23, 6.36_29]]  )
 elif model_name == "groupvit-gcc-redcaps":
  __a :  List[str]						=  torch.tensor([[16.18_73, 8.62_30]]  )
 else:
  raise ValueError(f"Model name {model_name} not supported."  )
 assert torch.allclose(outputs.logits_per_image  ,   lowerCAmelCase__  ,   atol=1e-3  )

 processor.save_pretrained(lowerCAmelCase__  )
 model.save_pretrained(lowerCAmelCase__  )
 print('''Successfully saved processor and model to'''  ,   lowerCAmelCase__  )

 if push_to_hub:
  print('''Pushing to the hub...'''  )
  processor.push_to_hub(lowerCAmelCase__  ,   organization='''nielsr'''  )
  model.push_to_hub(lowerCAmelCase__  ,   organization='''nielsr'''  )


if __name__ == "__main__":
      lowercase__       =argparse.ArgumentParser()
      parser.add_argument(
          '--pytorch_dump_folder_path', default=None, type=str, help='Path to dump the processor and PyTorch model.'
      )
      parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to GroupViT checkpoint')
      parser.add_argument(
          '--model_name',
          default='groupvit-gccy-fcc',
          type=str,
          help='Name of the model. Expecting either \'groupvit-gcc-yfcc\' or \'groupvit-gcc-redcaps\'',
      )
      parser.add_argument(
          '--push_to_hub',
          action='store_true',
          help='Whether or not to push the converted model and processor to the 🤗 hub using the provided `model_name`.',
      )
      lowercase__       =parser.parse_args()

      convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)
