


import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


__UpperCamelCase					:	List[Any]          =				{
    'vit_b32': {
        'embed_dim': 512,
        'image_resolution': 768,
        'context_length': 16,
        'vocab_size': 4_9408,
        'vision_layers': 12,
        'vision_width': 768,
        'vision_patch_size': 32,
        'transformer_width': 512,
        'transformer_heads': 8,
        'transformer_layers': 12,
    },
    'vit_b16': {
        'embed_dim': 512,
        'image_resolution': 768,
        'context_length': 16,
        'vocab_size': 4_9408,
        'vision_layers': 12,
        'vision_width': 768,
        'vision_patch_size': 16,
        'transformer_width': 512,
        'transformer_heads': 8,
        'transformer_layers': 12,
    },
    'vit_l14': {
        'embed_dim': 768,
        'image_resolution': 840,
        'context_length': 16,
        'vocab_size': 4_9408,
        'vision_layers': 24,
        'vision_width': 1024,
        'vision_patch_size': 14,
        'transformer_width': 768,
        'transformer_heads': 12,
        'transformer_layers': 12,
    },
}
def   snake_case_    (			__lowercase      ,							__lowercase=""      ,							__lowercase="/"							):
   UpperCAmelCase_				:   List[str]							= []

   for k, v in params.items():
      UpperCAmelCase_				:   Dict							= parent_key + sep + k if parent_key else k

      if isinstance(__lowercase      ,							collections.MutableMapping							):
         items.extend(flatten_nested_dict(__lowercase      ,							__lowercase      ,							sep=__lowercase							).items()							)
      else:
         items.append((new_key, v)							)
   return dict(__lowercase							)
def   snake_case_    (			__lowercase							):
   return jax.tree_util.tree_map(lambda __lowercase							: x.astype(jnp.floataa							) if x.dtype == jnp.bfloataa else x      ,							__lowercase							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_				:   Optional[int]							= pt_attn_layer.in_proj_weight.chunk(3      ,							dim=0							)
   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_				:   Tuple							= pt_attn_layer.in_proj_bias.chunk(3      ,							dim=0							)

   UpperCAmelCase_				:   Dict							= pt_attn_layer.out_proj.weight
   UpperCAmelCase_				:   List[str]							= pt_attn_layer.out_proj.bias

   UpperCAmelCase_				:   Tuple							= q_proj
   UpperCAmelCase_				:   Optional[int]							= q_proj_bias

   UpperCAmelCase_				:   List[Any]							= k_proj
   UpperCAmelCase_				:   int							= k_proj_bias

   UpperCAmelCase_				:   Tuple							= v_proj
   UpperCAmelCase_				:   int							= v_proj_bias

   UpperCAmelCase_				:   Tuple							= out_proj_weights
   UpperCAmelCase_				:   List[Any]							= out_proj_bias
def   snake_case_    (			__lowercase      ,							__lowercase							):
   copy_linear(hf_mlp.fca      ,							pt_mlp.c_fc							)
   copy_linear(hf_mlp.fca      ,							pt_mlp.c_proj							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   UpperCAmelCase_				:   int							= pt_linear.weight
   UpperCAmelCase_				:   List[str]							= pt_linear.bias
def   snake_case_    (			__lowercase      ,							__lowercase							):
   # copy layer norms
   copy_linear(hf_layer.layer_norma      ,							pt_layer.ln_a							)
   copy_linear(hf_layer.layer_norma      ,							pt_layer.ln_a							)

   # copy MLP
   copy_mlp(hf_layer.mlp      ,							pt_layer.mlp							)

   # copy attn
   copy_attn_layer(hf_layer.self_attn      ,							pt_layer.attn							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   for hf_layer, pt_layer in zip(__lowercase      ,							__lowercase							):
      copy_layer(__lowercase      ,							__lowercase							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   # copy  embeds
   UpperCAmelCase_				:   Union[str, Any]							= pt_model.token_embedding.weight
   UpperCAmelCase_				:   Union[str, Any]							= pt_model.positional_embedding

   # copy layer norm
   copy_linear(hf_encoder.final_layer_norm      ,							pt_model.ln_final							)

   # copy hidden layers
   copy_layers(hf_encoder.encoder.layers      ,							pt_model.transformer.resblocks							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   # copy projection
   UpperCAmelCase_				:   Optional[Any]							= pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model      ,							__lowercase							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   # copy projection
   UpperCAmelCase_				:   Optional[Any]							= pt_model.visual.proj.data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layernorm      ,							pt_model.visual.ln_pre							)
   copy_linear(hf_model.vision_model.post_layernorm      ,							pt_model.visual.ln_post							)

   # copy embeds
   UpperCAmelCase_				:   List[Any]							= pt_model.visual.conva.weight.data
   UpperCAmelCase_				:   Optional[Any]							= pt_model.visual.class_embedding
   UpperCAmelCase_				:   int							= pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers      ,							pt_model.visual.transformer.resblocks							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   UpperCAmelCase_				:   List[str]							= flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']							)

   UpperCAmelCase_				:   str							= torch.from_numpy(flax_class_token_params['''scale''']							)
   UpperCAmelCase_				:   Union[str, Any]							= torch.from_numpy(flax_class_token_params['''bias''']							)
   UpperCAmelCase_				:   List[Any]							= nn.Parameter(__lowercase							)
   UpperCAmelCase_				:   List[str]							= nn.Parameter(__lowercase							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   UpperCAmelCase_				:   int							= hf_model.state_dict()
   UpperCAmelCase_				:   Any							= {}

   # Rename class prediction head flax params to pytorch HF
   UpperCAmelCase_				:   Any							= flatten_nested_dict(flax_params['''class_head''']							)

   for flax_key, v in flax_class_params.items():
      UpperCAmelCase_				:   Any							= flax_key.replace('''/'''      ,							'''.'''							)
      UpperCAmelCase_				:   str							= torch_key.replace('''.kernel'''      ,							'''.weight'''							)
      UpperCAmelCase_				:   Dict							= torch_key.replace('''Dense_0'''      ,							'''dense0'''							)
      UpperCAmelCase_				:   List[str]							= '''class_head.''' + torch_key

      if "weight" in torch_key and v.ndim == 2:
         UpperCAmelCase_				:   Any							= v.T

      UpperCAmelCase_				:   Optional[int]							= nn.Parameter(torch.from_numpy(__lowercase							)							)

   # Rename box prediction box flax params to pytorch HF
   UpperCAmelCase_				:   Optional[int]							= flatten_nested_dict(flax_params['''obj_box_head''']							)

   for flax_key, v in flax_box_params.items():
      UpperCAmelCase_				:   str							= flax_key.replace('''/'''      ,							'''.'''							)
      UpperCAmelCase_				:   int							= torch_key.replace('''.kernel'''      ,							'''.weight'''							)
      UpperCAmelCase_				:   Dict							= torch_key.replace('''_'''      ,							''''''							).lower()
      UpperCAmelCase_				:   List[str]							= '''box_head.''' + torch_key

      if "weight" in torch_key and v.ndim == 2:
         UpperCAmelCase_				:   Optional[int]							= v.T

      UpperCAmelCase_				:   str							= nn.Parameter(torch.from_numpy(__lowercase							)							)

   # Copy flax params to PyTorch params
   for name, param in new_params.items():
      if name in pt_params.keys():
         pt_params[name].copy_(__lowercase							)
def   snake_case_    (			__lowercase      ,							__lowercase							):
   for k, v in flax_attn_params.items():
      if k.startswith('''transformer'''							):
         UpperCAmelCase_				:   List[str]							= k.replace('''transformer.resblocks'''      ,							'''text_model.encoder.layers'''							)
      else:
         UpperCAmelCase_				:   Tuple							= k.replace('''visual.transformer.resblocks'''      ,							'''vision_model.encoder.layers'''							)

      UpperCAmelCase_				:   Tuple							= torch_key.replace('''attn'''      ,							'''self_attn'''							)
      UpperCAmelCase_				:   Optional[int]							= torch_key.replace('''key'''      ,							'''k_proj'''							)
      UpperCAmelCase_				:   Optional[int]							= torch_key.replace('''value'''      ,							'''v_proj'''							)
      UpperCAmelCase_				:   List[Any]							= torch_key.replace('''query'''      ,							'''q_proj'''							)
      UpperCAmelCase_				:   List[str]							= torch_key.replace('''out'''      ,							'''out_proj'''							)

      if "bias" in torch_key and v.ndim == 2:
         UpperCAmelCase_				:   Optional[Any]							= v.shape[0] * v.shape[1]
         UpperCAmelCase_				:   List[Any]							= v.reshape(__lowercase							)

      if "weight" in torch_key and "out" in torch_key:
         UpperCAmelCase_				:   Union[str, Any]							= (v.shape[0] * v.shape[1], v.shape[2])
         UpperCAmelCase_				:   List[str]							= v.reshape(__lowercase							).T

      if "weight" in torch_key and "out" not in torch_key:
         UpperCAmelCase_				:   Any							= (v.shape[0], v.shape[1] * v.shape[2])
         UpperCAmelCase_				:   List[str]							= v.reshape(__lowercase							).T

      # Copy flax CLIP attn params to HF PyTorch params
      UpperCAmelCase_				:   Union[str, Any]							= torch.from_numpy(__lowercase							)
      hf_backbone.state_dict()[torch_key].copy_(__lowercase							)
def   snake_case_    (			__lowercase							):
   UpperCAmelCase_				:   List[str]							= {}
   UpperCAmelCase_				:   Union[str, Any]							= []

   for k, v in params.items():
      if "attn." in k:
         UpperCAmelCase_				:   List[Any]							= k[: k.rindex('''attn.'''							) + 5]
         if base in processed_attn_layers:
            continue

         processed_attn_layers.append(__lowercase							)
         UpperCAmelCase_				:   List[str]							= params[base + '''out.weight'''].shape[-1]
         UpperCAmelCase_				:   List[Any]							= params[base + '''out.weight'''].reshape(__lowercase      ,							__lowercase							).T
         UpperCAmelCase_				:   int							= params[base + '''out.bias''']
      else:
         UpperCAmelCase_				:   Any							= v
   return new_params
def   snake_case_    (			__lowercase      ,							__lowercase							):
   UpperCAmelCase_				:   Optional[int]							= CLIP(**__lowercase							)
   torch_model.eval()
   UpperCAmelCase_				:   Optional[int]							= torch_model.state_dict()

   UpperCAmelCase_				:   Any							= flatten_nested_dict(flax_params['''backbone''']['''clip''']							)
   UpperCAmelCase_				:   Any							= {}

   for flax_key, v in flax_clip_params.items():
      UpperCAmelCase_				:   Optional[Any]							= flax_key.replace('''/'''      ,							'''.'''							)
      UpperCAmelCase_				:   Tuple							= torch_key.replace('''text.token_embedding.embedding'''      ,							'''token_embedding.kernel'''							)

      if (
          torch_key.startswith('''text.transformer'''							)
          or torch_key.startswith('''text.text_projection'''							)
          or torch_key.startswith('''text.ln_final'''							)
          or torch_key.startswith('''text.positional_embedding'''							)
      ):
         UpperCAmelCase_				:   List[Any]							= torch_key[5:]

      UpperCAmelCase_				:   Union[str, Any]							= torch_key.replace('''text_projection.kernel'''      ,							'''text_projection'''							)
      UpperCAmelCase_				:   Dict							= torch_key.replace('''visual.proj.kernel'''      ,							'''visual.proj'''							)
      UpperCAmelCase_				:   List[str]							= torch_key.replace('''.scale'''      ,							'''.weight'''							)
      UpperCAmelCase_				:   List[str]							= torch_key.replace('''.kernel'''      ,							'''.weight'''							)

      if "conv" in torch_key or "downsample.0.weight" in torch_key:
         UpperCAmelCase_				:   Optional[Any]							= v.transpose(3      ,							2      ,							0      ,							1							)

      elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
         # Fully connected layers are transposed, embeddings are not
         UpperCAmelCase_				:   Dict							= v.T

      UpperCAmelCase_				:   str							= v

   UpperCAmelCase_				:   Optional[Any]							= _convert_attn_layers(__lowercase							)
   new_torch_params.update(__lowercase							)
   UpperCAmelCase_				:   Optional[int]							= {}

   # Copy flax CLIP backbone params to PyTorch params
   for name, param in new_torch_params.items():
      if name in torch_clip_params.keys():
         UpperCAmelCase_				:   Any							= torch.from_numpy(new_torch_params[name]							)
         torch_clip_params[name].copy_(__lowercase							)
      else:
         UpperCAmelCase_				:   Optional[int]							= param

   return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase=None							):
   UpperCAmelCase_				:   List[str]							= Repository(__lowercase      ,							clone_from=F'''google/{pytorch_dump_folder_path}'''							)
   repo.git_pull()

   if config_path is not None:
      UpperCAmelCase_				:   List[str]							= OwlViTConfig.from_pretrained(__lowercase							)
   else:
      UpperCAmelCase_				:   List[Any]							= OwlViTConfig()

   UpperCAmelCase_				:   Optional[int]							= OwlViTModel(__lowercase							).eval()
   UpperCAmelCase_				:   Tuple							= OwlViTForObjectDetection(__lowercase							).eval()

   copy_text_model_and_projection(__lowercase      ,							__lowercase							)
   copy_vision_model_and_projection(__lowercase      ,							__lowercase							)
   UpperCAmelCase_				:   List[Any]							= pt_backbone.logit_scale
   copy_flax_attn_params(__lowercase      ,							__lowercase							)

   UpperCAmelCase_				:   int							= hf_backbone
   copy_class_merge_token(__lowercase      ,							__lowercase							)
   copy_class_box_heads(__lowercase      ,							__lowercase							)

   # Save HF model
   hf_model.save_pretrained(repo.local_dir							)

   # Initialize image processor
   UpperCAmelCase_				:   Any							= OwlViTImageProcessor(
       size=config.vision_config.image_size      ,							crop_size=config.vision_config.image_size							)
   # Initialize tokenizer
   UpperCAmelCase_				:   str							= CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''      ,							pad_token='''!'''      ,							model_max_length=1_6							)

   # Initialize processor
   UpperCAmelCase_				:   str							= OwlViTProcessor(image_processor=__lowercase      ,							tokenizer=__lowercase							)
   image_processor.save_pretrained(repo.local_dir							)
   processor.save_pretrained(repo.local_dir							)

   repo.git_add()
   repo.git_commit('''Upload model and processor'''							)
   repo.git_push()


if __name__ == "__main__":
       __UpperCamelCase					:	Optional[Any]          =				argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--owlvit_version',
           default=None,
           type=str,
           required=True,
           help='OWL-ViT model name [clip_b16, clip_b32, clip_l14].',
       )
       parser.add_argument(
           '--owlvit_checkpoint', default=None, type=str, required=True, help='Path to flax model checkpoint.'
       )
       parser.add_argument('--hf_config', default=None, type=str, required=True, help='Path to HF model config.')
       parser.add_argument(
           '--pytorch_dump_folder_path', default='hf_model', type=str, help='Path to the output PyTorch model.'
       )
       __UpperCamelCase					:	Tuple          =				parser.parse_args()

       # Initialize PyToch clip model
       __UpperCamelCase					:	Any          =				args.owlvit_version
       if model_name == "clip_b16":
              __UpperCamelCase					:	str          =				CONFIGS['vit_b16']
       elif model_name == "clip_b32":
              __UpperCamelCase					:	Optional[Any]          =				CONFIGS['vit_b32']
       elif model_name == "clip_l14":
              __UpperCamelCase					:	Dict          =				CONFIGS['vit_l14']

       # Load from checkpoint and convert params to float-32
       __UpperCamelCase					:	Any          =				checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['optimizer']['target']
       __UpperCamelCase					:	Any          =				jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
       del variables

       # Convert CLIP backbone
       __UpperCamelCase       ,       __UpperCamelCase       ,       __UpperCamelCase					:	Optional[Any]          =				convert_clip_backbone(flax_params, torch_config)

       convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)



import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       from transformers.models.electra.modeling_flax_electra import (
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForMultipleChoice,
           FlaxElectraForPreTraining,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForSequenceClassification,
           FlaxElectraForTokenClassification,
           FlaxElectraModel,
       )





class 			lowerCAmelCase__(     unittest.TestCase    ):






    '''simple docstring'''





    def __init__(						self					: Any							, __snake_case					: Tuple							, __snake_case					: Union[str, Any]=13							, __snake_case					: Tuple=7							, __snake_case					: str=True							, __snake_case					: List[Any]=True							, __snake_case					: Optional[int]=True							, __snake_case					: List[str]=True							, __snake_case					: Tuple=99							, __snake_case					: int=24							, __snake_case					: int=32							, __snake_case					: Tuple=5							, __snake_case					: int=4							, __snake_case					: Any=37							, __snake_case					: int="gelu"							, __snake_case					: Union[str, Any]=0.1							, __snake_case					: Optional[Any]=0.1							, __snake_case					: str=512							, __snake_case					: List[Any]=16							, __snake_case					: Tuple=2							, __snake_case					: List[str]=0.02							, __snake_case					: Any=4							, ):





       '''simple docstring'''
       UpperCAmelCase_				:   Union[str, Any]							= parent
       UpperCAmelCase_				:   Dict							= batch_size
       UpperCAmelCase_				:   Union[str, Any]							= seq_length
       UpperCAmelCase_				:   Any							= is_training
       UpperCAmelCase_				:   Tuple							= use_attention_mask
       UpperCAmelCase_				:   Dict							= use_token_type_ids
       UpperCAmelCase_				:   str							= use_labels
       UpperCAmelCase_				:   Union[str, Any]							= vocab_size
       UpperCAmelCase_				:   Union[str, Any]							= hidden_size
       UpperCAmelCase_				:   Union[str, Any]							= embedding_size
       UpperCAmelCase_				:   Optional[int]							= num_hidden_layers
       UpperCAmelCase_				:   Tuple							= num_attention_heads
       UpperCAmelCase_				:   List[Any]							= intermediate_size
       UpperCAmelCase_				:   Union[str, Any]							= hidden_act
       UpperCAmelCase_				:   Any							= hidden_dropout_prob
       UpperCAmelCase_				:   List[str]							= attention_probs_dropout_prob
       UpperCAmelCase_				:   List[str]							= max_position_embeddings
       UpperCAmelCase_				:   Dict							= type_vocab_size
       UpperCAmelCase_				:   Tuple							= type_sequence_label_size
       UpperCAmelCase_				:   Dict							= initializer_range
       UpperCAmelCase_				:   Any							= num_choices





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Any							= ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size						)

       UpperCAmelCase_				:   Any							= None
       if self.use_attention_mask:
          UpperCAmelCase_				:   List[str]							= random_attention_mask([self.batch_size, self.seq_length]						)

       UpperCAmelCase_				:   List[str]							= None
       if self.use_token_type_ids:
          UpperCAmelCase_				:   str							= ids_tensor([self.batch_size, self.seq_length]							, self.type_vocab_size						)

       UpperCAmelCase_				:   int							= ElectraConfig(
           vocab_size=self.vocab_size							, hidden_size=self.hidden_size							, embedding_size=self.embedding_size							, num_hidden_layers=self.num_hidden_layers							, num_attention_heads=self.num_attention_heads							, intermediate_size=self.intermediate_size							, hidden_act=self.hidden_act							, hidden_dropout_prob=self.hidden_dropout_prob							, attention_probs_dropout_prob=self.attention_probs_dropout_prob							, max_position_embeddings=self.max_position_embeddings							, type_vocab_size=self.type_vocab_size							, initializer_range=self.initializer_range							, )

       return config, input_ids, token_type_ids, attention_mask





    def        _lowerCamelCase			(						self					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= self.prepare_config_and_inputs()
       UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_				:   Union[str, Any]							= config_and_inputs
       UpperCAmelCase_				:   Optional[Any]							= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict





@require_flax
class 			lowerCAmelCase__(     snake_case__     ,     unittest.TestCase    ):






    '''simple docstring'''





    A_	:  Any								=			True

    A_	:  int								=			(
        (
            FlaxElectraModel,
            FlaxElectraForCausalLM,
            FlaxElectraForMaskedLM,
            FlaxElectraForPreTraining,
            FlaxElectraForTokenClassification,
            FlaxElectraForQuestionAnswering,
            FlaxElectraForMultipleChoice,
            FlaxElectraForSequenceClassification,
        )
        if is_flax_available()
        else ()
    )





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= FlaxElectraModelTester(self						)





    @slow
    def        _lowerCamelCase			(						self					: Union[str, Any]						):





       '''simple docstring'''
       for model_class_name in self.all_model_classes:
          if model_class_name == FlaxElectraForMaskedLM:
             UpperCAmelCase_				:   List[str]							= model_class_name.from_pretrained('''google/electra-small-generator'''						)
          else:
             UpperCAmelCase_				:   Optional[Any]							= model_class_name.from_pretrained('''google/electra-small-discriminator'''						)
          UpperCAmelCase_				:   Tuple							= model(np.ones((1, 1)						)						)
          self.assertIsNotNone(__snake_case						)