


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,
)


__lowerCAmelCase									=				{
    '''vit_b32''': {
        '''embed_dim''': 5_12,
        '''image_resolution''': 7_68,
        '''context_length''': 16,
        '''vocab_size''': 4_94_08,
        '''vision_layers''': 12,
        '''vision_width''': 7_68,
        '''vision_patch_size''': 32,
        '''transformer_width''': 5_12,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_b16''': {
        '''embed_dim''': 5_12,
        '''image_resolution''': 7_68,
        '''context_length''': 16,
        '''vocab_size''': 4_94_08,
        '''vision_layers''': 12,
        '''vision_width''': 7_68,
        '''vision_patch_size''': 16,
        '''transformer_width''': 5_12,
        '''transformer_heads''': 8,
        '''transformer_layers''': 12,
    },
    '''vit_l14''': {
        '''embed_dim''': 7_68,
        '''image_resolution''': 8_40,
        '''context_length''': 16,
        '''vocab_size''': 4_94_08,
        '''vision_layers''': 24,
        '''vision_width''': 10_24,
        '''vision_patch_size''': 14,
        '''transformer_width''': 7_68,
        '''transformer_heads''': 12,
        '''transformer_layers''': 12,
    },
}




def   snake_case_	(    snake_case     ,				snake_case=""     ,				snake_case="/"	)    ->							Union[str, Any]:
       lowercase__:  List[str]         =  []

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

              if isinstance(snake_case     ,				collections.MutableMapping	):
                     items.extend(flatten_nested_dict(snake_case     ,				snake_case     ,				sep=snake_case	).items()	)
              else:
                     items.append((new_key, v)	)
       return dict(snake_case	)




def   snake_case_	(    snake_case	)    ->							Optional[int]:
       return jax.tree_util.tree_map(lambda snake_case	: x.astype(jnp.floataa	) if x.dtype == jnp.bfloataa else x     ,				snake_case	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							Dict:
       lowercase__     ,		lowercase__     ,		lowercase__:  int         =  pt_attn_layer.in_proj_weight.chunk(3     ,				dim=0	)
       lowercase__     ,		lowercase__     ,		lowercase__:  Optional[int]         =  pt_attn_layer.in_proj_bias.chunk(3     ,				dim=0	)

       lowercase__:  Optional[int]         =  pt_attn_layer.out_proj.weight
       lowercase__:  str         =  pt_attn_layer.out_proj.bias

       lowercase__:  Optional[Any]         =  q_proj
       lowercase__:  Any         =  q_proj_bias

       lowercase__:  Optional[int]         =  k_proj
       lowercase__:  Union[str, Any]         =  k_proj_bias

       lowercase__:  List[Any]         =  v_proj
       lowercase__:  List[str]         =  v_proj_bias

       lowercase__:  int         =  out_proj_weights
       lowercase__:  Any         =  out_proj_bias




def   snake_case_	(    snake_case     ,				snake_case	)    ->							Dict:
       copy_linear(hf_mlp.fca     ,				pt_mlp.c_fc	)
       copy_linear(hf_mlp.fca     ,				pt_mlp.c_proj	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							int:
       lowercase__:  List[str]         =  pt_linear.weight
       lowercase__:  Any         =  pt_linear.bias




def   snake_case_	(    snake_case     ,				snake_case	)    ->							str:
       # 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_	(    snake_case     ,				snake_case	)    ->							Tuple:
       for hf_layer, pt_layer in zip(snake_case     ,				snake_case	):
              copy_layer(snake_case     ,				snake_case	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							int:
       # copy  embeds
       lowercase__:  Dict         =  pt_model.token_embedding.weight
       lowercase__:  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_	(    snake_case     ,				snake_case	)    ->							int:
       # copy projection
       lowercase__:  List[Any]         =  pt_model.text_projection.data.T

       # copy text encoder
       copy_encoder(hf_model.text_model     ,				snake_case	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							List[Any]:
       # copy projection
       lowercase__:  Dict         =  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
       lowercase__:  List[Any]         =  pt_model.visual.conva.weight.data
       lowercase__:  str         =  pt_model.visual.class_embedding
       lowercase__:  Optional[Any]         =  pt_model.visual.positional_embedding.data

       # copy encoder
       copy_layers(hf_model.vision_model.encoder.layers     ,				pt_model.visual.transformer.resblocks	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							Any:
       lowercase__:  Optional[int]         =  flatten_nested_dict(flax_params['backbone']['merged_class_token']	)

       lowercase__:  str         =  torch.from_numpy(flax_class_token_params['scale']	)
       lowercase__:  str         =  torch.from_numpy(flax_class_token_params['bias']	)
       lowercase__:  str         =  nn.Parameter(snake_case	)
       lowercase__:  Any         =  nn.Parameter(snake_case	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							int:
       lowercase__:  Tuple         =  hf_model.state_dict()
       lowercase__:  Optional[Any]         =  {}

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

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

              if "weight" in torch_key and v.ndim == 2:
                     lowercase__:  List[Any]         =  v.T

              lowercase__:  Tuple         =  nn.Parameter(torch.from_numpy(snake_case	)	)

       # Rename box prediction box flax params to pytorch HF
       lowercase__:  Union[str, Any]         =  flatten_nested_dict(flax_params['obj_box_head']	)

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

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

              lowercase__:  str         =  nn.Parameter(torch.from_numpy(snake_case	)	)

       # Copy flax params to PyTorch params
       for name, param in new_params.items():
              if name in pt_params.keys():
                     pt_params[name].copy_(snake_case	)




def   snake_case_	(    snake_case     ,				snake_case	)    ->							Any:
       for k, v in flax_attn_params.items():
              if k.startswith('transformer'	):
                     lowercase__:  str         =  k.replace('transformer.resblocks'     ,				'text_model.encoder.layers'	)
              else:
                     lowercase__:  Union[str, Any]         =  k.replace('visual.transformer.resblocks'     ,				'vision_model.encoder.layers'	)

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

              if "bias" in torch_key and v.ndim == 2:
                     lowercase__:  List[str]         =  v.shape[0] * v.shape[1]
                     lowercase__:  List[Any]         =  v.reshape(snake_case	)

              if "weight" in torch_key and "out" in torch_key:
                     lowercase__:  Optional[int]         =  (v.shape[0] * v.shape[1], v.shape[2])
                     lowercase__:  Tuple         =  v.reshape(snake_case	).T

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

              # Copy flax CLIP attn params to HF PyTorch params
              lowercase__:  Any         =  torch.from_numpy(snake_case	)
              hf_backbone.state_dict()[torch_key].copy_(snake_case	)




def   snake_case_	(    snake_case	)    ->							List[Any]:
       lowercase__:  Optional[Any]         =  {}
       lowercase__:  List[Any]         =  []

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

                     processed_attn_layers.append(snake_case	)
                     lowercase__:  List[str]         =  params[base + 'out.weight'].shape[-1]
                     lowercase__:  str         =  params[base + 'out.weight'].reshape(snake_case     ,				snake_case	).T
                     lowercase__:  int         =  params[base + 'out.bias']
              else:
                     lowercase__:  List[str]         =  v
       return new_params




def   snake_case_	(    snake_case     ,				snake_case	)    ->							Any:
       lowercase__:  str         =  CLIP(**snake_case	)
       torch_model.eval()
       lowercase__:  Any         =  torch_model.state_dict()

       lowercase__:  Dict         =  flatten_nested_dict(flax_params['backbone']['clip']	)
       lowercase__:  Tuple         =  {}

       for flax_key, v in flax_clip_params.items():
              lowercase__:  Optional[int]         =  flax_key.replace('/'     ,				'.'	)
              lowercase__:  Any         =  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'	)
              ):
                     lowercase__:  Dict         =  torch_key[5:]

              lowercase__:  Any         =  torch_key.replace('text_projection.kernel'     ,				'text_projection'	)
              lowercase__:  Tuple         =  torch_key.replace('visual.proj.kernel'     ,				'visual.proj'	)
              lowercase__:  Optional[int]         =  torch_key.replace('.scale'     ,				'.weight'	)
              lowercase__:  List[str]         =  torch_key.replace('.kernel'     ,				'.weight'	)

              if "conv" in torch_key or "downsample.0.weight" in torch_key:
                     lowercase__:  Dict         =  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
                     lowercase__:  List[Any]         =  v.T

              lowercase__:  str         =  v

       lowercase__:  str         =  _convert_attn_layers(snake_case	)
       new_torch_params.update(snake_case	)
       lowercase__:  str         =  {}

       # Copy flax CLIP backbone params to PyTorch params
       for name, param in new_torch_params.items():
              if name in torch_clip_params.keys():
                     lowercase__:  Union[str, Any]         =  torch.from_numpy(new_torch_params[name]	)
                     torch_clip_params[name].copy_(snake_case	)
              else:
                     lowercase__:  Union[str, Any]         =  param

       return torch_clip_params, torch_model, attn_params




@torch.no_grad()
def   snake_case_	(    snake_case     ,				snake_case     ,				snake_case     ,				snake_case     ,				snake_case=None	)    ->							int:
       lowercase__:  Tuple         =  Repository(snake_case     ,				clone_from=f'google/{pytorch_dump_folder_path}'	)
       repo.git_pull()

       if config_path is not None:
              lowercase__:  Tuple         =  OwlViTConfig.from_pretrained(snake_case	)
       else:
              lowercase__:  str         =  OwlViTConfig()

       lowercase__:  Union[str, Any]         =  OwlViTModel(snake_case	).eval()
       lowercase__:  Tuple         =  OwlViTForObjectDetection(snake_case	).eval()

       copy_text_model_and_projection(snake_case     ,				snake_case	)
       copy_vision_model_and_projection(snake_case     ,				snake_case	)
       lowercase__:  Optional[int]         =  pt_backbone.logit_scale
       copy_flax_attn_params(snake_case     ,				snake_case	)

       lowercase__:  List[str]         =  hf_backbone
       copy_class_merge_token(snake_case     ,				snake_case	)
       copy_class_box_heads(snake_case     ,				snake_case	)

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

       # Initialize image processor
       lowercase__:  Tuple         =  OwlViTImageProcessor(
           size=config.vision_config.image_size     ,				crop_size=config.vision_config.image_size	)
       # Initialize tokenizer
       lowercase__:  Optional[Any]         =  CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'     ,				pad_token='!'     ,				model_max_length=16	)

       # Initialize processor
       lowercase__:  int         =  OwlViTProcessor(image_processor=snake_case     ,				tokenizer=snake_case	)
       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__":
 __lowerCAmelCase									=				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.'''
 )
 __lowerCAmelCase									=				parser.parse_args()

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

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

 # Convert CLIP backbone
 __lowerCAmelCase			,__lowerCAmelCase			,__lowerCAmelCase									=				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

from cer import CER


__lowerCAmelCase									=				CER()





class     __a     (       unittest.TestCase ):
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[Any]:




        '''simple docstring'''
        lowercase__:  Optional[Any]         =  ['White House']
        lowercase__:  Optional[int]         =  ['white house']
        # S = 2, D = 0, I = 0, N = 11, CER = 2 / 11
        lowercase__:  str         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.1_8_1_8_1_8_1_8_1_8					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							List[str]:




        '''simple docstring'''
        lowercase__:  Union[str, Any]         =  ['were wolf']
        lowercase__:  Union[str, Any]         =  ['werewolf']
        # S = 0, D = 0, I = 1, N = 9, CER = 1 / 9
        lowercase__:  Tuple         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.1_1_1_1_1_1_1					) < 1E-6					)

        lowercase__:  Dict         =  ['werewolf']
        lowercase__:  Optional[Any]         =  ['weae     wolf']
        # S = 1, D = 1, I = 0, N = 8, CER = 0.25
        lowercase__:  Optional[int]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.2_5					) < 1E-6					)

        # consecutive whitespaces case 1
        lowercase__:  Union[str, Any]         =  ['were wolf']
        lowercase__:  int         =  ['were               wolf']
        # S = 0, D = 0, I = 0, N = 9, CER = 0
        lowercase__:  int         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.0					) < 1E-6					)

        # consecutive whitespaces case 2
        lowercase__:  List[str]         =  ['were   wolf']
        lowercase__:  Tuple         =  ['were               wolf']
        # S = 0, D = 0, I = 0, N = 9, CER = 0
        lowercase__:  Dict         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.0					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							List[str]:




        '''simple docstring'''
        lowercase__:  Optional[int]         =  ['werewolf']
        lowercase__:  Dict         =  ['weaewolf']
        # S = 1, D = 0, I = 0, N = 8, CER = 0.125
        lowercase__:  int         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.1_2_5					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[Any]:




        '''simple docstring'''
        lowercase__:  int         =  ['werewolf']
        lowercase__:  Optional[int]         =  ['wereawolf']
        # S = 0, D = 1, I = 0, N = 8, CER = 0.125
        lowercase__:  Dict         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.1_2_5					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[Any]:




        '''simple docstring'''
        lowercase__:  Optional[Any]         =  ['werewolf']
        lowercase__:  List[str]         =  ['wereolf']
        # S = 0, D = 0, I = 1, N = 8, CER = 0.125
        lowercase__:  Optional[Any]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.1_2_5					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[int]:




        '''simple docstring'''
        lowercase__:  List[str]         =  ['werewolf']
        lowercase__:  List[Any]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertEqual(lowerCAmelCase__							,					0.0					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							List[Any]:




        '''simple docstring'''
        lowercase__:  int         =  ['werewolf', 'I am your father']
        lowercase__:  Union[str, Any]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertEqual(lowerCAmelCase__							,					0.0					)

        lowercase__:  List[str]         =  ['werewolf', 'I am your father', 'doge']
        lowercase__:  Any         =  ['werxwolf', 'I       am your father', 'doge']
        # S = 1, D = 0, I = 0, N = 28, CER = 1 / 28
        lowercase__:  int         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.0_3_5_7_1_4_2_8					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[int]:




        '''simple docstring'''
        lowercase__:  List[Any]         =  ['My hovercraft', 'is full of eels']
        lowercase__:  Union[str, Any]         =  ['My hovercraft is full', ' of eels']
        # S = 0, D = 0, I = 2, N = 28, CER = 2 / 28
        # whitespace at the front of " of eels" will be strip during preporcessing
        # so need to insert 2 whitespaces
        lowercase__:  Optional[int]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__							,					concatenate_texts=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.0_7_1_4_2_8					) < 1E-6					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Dict:




        '''simple docstring'''
        lowercase__:  Dict         =  ['我能吞下玻璃而不伤身体']
        lowercase__:  Tuple         =  [' 能吞虾玻璃而 不霜身体啦']
        # S = 3, D = 2, I = 0, N = 11, CER = 5 / 11
        lowercase__:  str         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5					) < 1E-6					)

        lowercase__:  Any         =  ['我能吞下玻璃', '而不伤身体']
        lowercase__:  Union[str, Any]         =  ['我    能 吞 下 玻 璃', '而不伤身体']
        # S = 0, D = 5, I = 0, N = 11, CER = 5 / 11
        lowercase__:  Optional[Any]         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5_4_5					) < 1E-6					)

        lowercase__:  int         =  ['我能吞下玻璃而不伤身体']
        lowercase__:  Any         =  cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)
        self.assertFalse(lowerCAmelCase__							,					0.0					)
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							List[str]:




        '''simple docstring'''
        lowercase__:  Dict         =  ['']
        lowercase__:  int         =  ['Hypothesis']
        with self.assertRaises(lowerCAmelCase__					):
               cer.compute(predictions=lowerCAmelCase__							,					references=lowerCAmelCase__					)


if __name__ == "__main__":
 unittest.main()


