


'''simple docstring'''

import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel

def _SCREAMING_SNAKE_CASE (A			)			->			Tuple:






       """simple docstring"""







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

       return name

def _SCREAMING_SNAKE_CASE (A    ,					A			)			->			str:






       """simple docstring"""







       for key in orig_state_dict.copy().keys():
              lowercase__      =					orig_state_dict.pop(A			)

              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
                     lowercase__      =					key.split('''.'''			)
                     lowercase__     ,lowercase__      =					int(key_split[2]			), int(key_split[4]			)
                     lowercase__      =					config.vision_config.hidden_size
                     if "weight" in key:
                            lowercase__      =					val[:dim, :]
                            lowercase__      =					val[dim : dim * 2, :]
                            lowercase__      =					val[-dim:, :]
                     else:
                            lowercase__      =					val[:dim]
                            lowercase__      =					val[dim : dim * 2]
                            lowercase__      =					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
                     lowercase__      =					key.split('''.'''			)
                     lowercase__      =					int(key_split[3]			)
                     lowercase__      =					config.text_config.hidden_size
                     if "weight" in key:
                            lowercase__      =					val[:dim, :]
                            lowercase__      =					val[
                                dim : dim * 2, :
                            ]
                            lowercase__      =					val[-dim:, :]
                     else:
                            lowercase__      =					val[:dim]
                            lowercase__      =					val[dim : dim * 2]
                            lowercase__      =					val[-dim:]
              else:
                     lowercase__      =					rename_key(A			)
                     # 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
                     ):
                            lowercase__      =					val.squeeze_()
                     else:
                            lowercase__      =					val

       return orig_state_dict

def _SCREAMING_SNAKE_CASE ()			->			List[Any]:






       """simple docstring"""







       lowercase__      =					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
       lowercase__      =					Image.open(requests.get(A    ,					stream=A			).raw			)
       return im

@torch.no_grad()
def _SCREAMING_SNAKE_CASE (A    ,					A    ,					A="groupvit-gcc-yfcc"    ,					A=False			)			->			int:






       """simple docstring"""







       lowercase__      =					GroupViTConfig()
       lowercase__      =					GroupViTModel(A			).eval()

       lowercase__      =					torch.load(A    ,					map_location='''cpu'''			)['''model''']
       lowercase__      =					convert_state_dict(A    ,					A			)
       lowercase__     ,lowercase__      =					model.load_state_dict(A    ,					strict=A			)
       assert missing_keys == ["text_model.embeddings.position_ids"]
       assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(A			) == 0)

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

       with torch.no_grad():
              lowercase__      =					model(**A			)

       if model_name == "groupvit-gcc-yfcc":
              lowercase__      =					torch.tensor([[13.3_523, 6.3_629]]			)
       elif model_name == "groupvit-gcc-redcaps":
              lowercase__      =					torch.tensor([[16.1_873, 8.6_230]]			)
       else:
              raise ValueError(f"Model name {model_name} not supported."			)
       assert torch.allclose(outputs.logits_per_image    ,					A    ,					atol=1E-3			)

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

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


if __name__ == "__main__":
 lowerCamelCase				:  List[Any]    		=      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`.',
 )
 lowerCamelCase				:  List[Any]    		=      parser.parse_args()

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








'''simple docstring'''

from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule

def _SCREAMING_SNAKE_CASE (A    ,					A    ,					A=1E-12			)			->			str:






       """simple docstring"""







       lowercase__      =					jnp.divide(emb_a.T    ,					jnp.clip(jnp.linalg.norm(A    ,					axis=1			)    ,					a_min=A			)			).T
       lowercase__      =					jnp.divide(emb_a.T    ,					jnp.clip(jnp.linalg.norm(A    ,					axis=1			)    ,					a_min=A			)			).T
       return jnp.matmul(A    ,					norm_emb_a.T			)
class     __lowerCAmelCase		(nn.Module	):





      '''simple docstring'''






      lowerCAmelCase__   :	CLIPConfig
      lowerCAmelCase__   :	jnp.dtype       =				jnp.floataa



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             lowercase__      =					FlaxCLIPVisionModule(self.config.vision_config					)
             lowercase__      =					nn.Dense(self.config.projection_dim  ,		use_bias=UpperCamelCase  ,		dtype=self.dtype					)

             lowercase__      =					self.param('''concept_embeds'''  ,		jax.nn.initializers.ones  ,		(17, self.config.projection_dim)					)
             lowercase__      =					self.param(
                 '''special_care_embeds'''  ,		jax.nn.initializers.ones  ,		(3, self.config.projection_dim)					)

             lowercase__      =					self.param('''concept_embeds_weights'''  ,		jax.nn.initializers.ones  ,		(17,)					)
             lowercase__      =					self.param('''special_care_embeds_weights'''  ,		jax.nn.initializers.ones  ,		(3,)					)






      def __call__(self       :       Union[str, Any]  ,		UpperCamelCase       :       Tuple					):

             '''simple docstring'''







             lowercase__      =					self.vision_model(UpperCamelCase					)[1]
             lowercase__      =					self.visual_projection(UpperCamelCase					)

             lowercase__      =					jax_cosine_distance(UpperCamelCase  ,		self.special_care_embeds					)
             lowercase__      =					jax_cosine_distance(UpperCamelCase  ,		self.concept_embeds					)

             # increase this value to create a stronger `nfsw` filter
             # at the cost of increasing the possibility of filtering benign image inputs
             lowercase__      =					0.0

             lowercase__      =					special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
             lowercase__      =					jnp.round(UpperCamelCase  ,		3					)
             lowercase__      =					jnp.any(special_scores > 0  ,		axis=1  ,		keepdims=UpperCamelCase					)
             # Use a lower threshold if an image has any special care concept
             lowercase__      =					is_special_care * 0.01

             lowercase__      =					cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
             lowercase__      =					jnp.round(UpperCamelCase  ,		3					)
             lowercase__      =					jnp.any(concept_scores > 0  ,		axis=1					)

             return has_nsfw_concepts






class     __lowerCAmelCase		(lowercase_	):





      '''simple docstring'''






      lowerCAmelCase__   :	Union[str, Any]       =				CLIPConfig
      lowerCAmelCase__   :	Optional[int]       =				"""clip_input"""
      lowerCAmelCase__   :	Any       =				FlaxStableDiffusionSafetyCheckerModule



      def __init__(self       :       Optional[int]  ,		UpperCamelCase       :       CLIPConfig  ,		UpperCamelCase       :       Optional[Tuple] = None  ,		UpperCamelCase       :       int = 0  ,		UpperCamelCase       :       jnp.dtype = jnp.floataa  ,		UpperCamelCase       :       bool = True  ,		**UpperCamelCase       :       List[Any]  ,		):

             '''simple docstring'''







             if input_shape is None:
                    lowercase__      =					(1, 224, 224, 3)
             lowercase__      =					self.module_class(config=UpperCamelCase  ,		dtype=UpperCamelCase  ,		**UpperCamelCase					)
             super().__init__(UpperCamelCase  ,		UpperCamelCase  ,		input_shape=UpperCamelCase  ,		seed=UpperCamelCase  ,		dtype=UpperCamelCase  ,		_do_init=_do_init					)



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       jax.random.KeyArray  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       FrozenDict = None					):

             '''simple docstring'''







             lowercase__      =					jax.random.normal(UpperCamelCase  ,		UpperCamelCase					)

             lowercase__     ,lowercase__      =					jax.random.split(UpperCamelCase					)
             lowercase__      =					{'''params''': params_rng, '''dropout''': dropout_rng}

             lowercase__      =					self.module.init(UpperCamelCase  ,		UpperCamelCase					)['''params''']

             return random_params






      def __call__(self       :       Optional[Any]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       dict = None  ,		):

             '''simple docstring'''







             lowercase__      =					jnp.transpose(UpperCamelCase  ,		(0, 2, 3, 1)					)

             return self.module.apply(
                 {'''params''': params or self.params}  ,		jnp.array(UpperCamelCase  ,		dtype=jnp.floataa					)  ,		rngs={}  ,		)




