





'''simple docstring'''







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


A						    =							{
    '''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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	Dict=""       ,  lowerCAmelCase__     :	Dict="/")			->						List[Any]:

     '''simple docstring'''




     _lowercase							:      Tuple  						=							[]

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

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

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	str)			->						Dict:

     '''simple docstring'''




     return jax.tree_util.tree_map(lambda lowerCAmelCase__: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x       ,  lowerCAmelCase__)

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Any       ,  lowerCAmelCase__     :	Dict)			->						Optional[Any]:

     '''simple docstring'''




     _lowercase						, _lowercase						, _lowercase							:      List[Any]  						=							pt_attn_layer.in_proj_weight.chunk(3       ,  dim=0)
     _lowercase						, _lowercase						, _lowercase							:      Union[str, Any]  						=							pt_attn_layer.in_proj_bias.chunk(3       ,  dim=0)

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

     _lowercase							:      Any  						=							q_proj
     _lowercase							:      Tuple  						=							q_proj_bias

     _lowercase							:      Optional[Any]  						=							k_proj
     _lowercase							:      str  						=							k_proj_bias

     _lowercase							:      Any  						=							v_proj
     _lowercase							:      int  						=							v_proj_bias

     _lowercase							:      List[Any]  						=							out_proj_weights
     _lowercase							:      Optional[Any]  						=							out_proj_bias

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Dict       ,  lowerCAmelCase__     :	Tuple)			->						str:

     '''simple docstring'''




     copy_linear(hf_mlp.fca       ,  pt_mlp.c_fc)
     copy_linear(hf_mlp.fca       ,  pt_mlp.c_proj)

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	str       ,  lowerCAmelCase__     :	Any)			->						Any:

     '''simple docstring'''




     _lowercase							:      Any  						=							pt_linear.weight
     _lowercase							:      str  						=							pt_linear.bias

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	int       ,  lowerCAmelCase__     :	Optional[int])			->						str:

     '''simple docstring'''




     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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Any       ,  lowerCAmelCase__     :	Optional[Any])			->						Optional[Any]:

     '''simple docstring'''




     for hf_layer, pt_layer in zip(lowerCAmelCase__       ,  lowerCAmelCase__):
          copy_layer(lowerCAmelCase__       ,  lowerCAmelCase__)

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Dict       ,  lowerCAmelCase__     :	Optional[int])			->						Any:

     '''simple docstring'''




     _lowercase							:      Any  						=							pt_model.token_embedding.weight
     _lowercase							:      Tuple  						=							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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Tuple       ,  lowerCAmelCase__     :	Optional[int])			->						Optional[int]:

     '''simple docstring'''




     _lowercase							:      Any  						=							pt_model.text_projection.data.T

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

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[int]       ,  lowerCAmelCase__     :	List[str])			->						Tuple:

     '''simple docstring'''




     _lowercase							:      Optional[int]  						=							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							:      Dict  						=							pt_model.visual.conva.weight.data
     _lowercase							:      Union[str, Any]  						=							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 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	List[Any]       ,  lowerCAmelCase__     :	str)			->						str:

     '''simple docstring'''




     _lowercase							:      Union[str, Any]  						=							flatten_nested_dict(flax_params['backbone']['merged_class_token'])

     _lowercase							:      Union[str, Any]  						=							torch.from_numpy(flax_class_token_params['scale'])
     _lowercase							:      Optional[Any]  						=							torch.from_numpy(flax_class_token_params['bias'])
     _lowercase							:      Optional[Any]  						=							nn.Parameter(lowerCAmelCase__)
     _lowercase							:      Dict  						=							nn.Parameter(lowerCAmelCase__)

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Dict       ,  lowerCAmelCase__     :	int)			->						Dict:

     '''simple docstring'''




     _lowercase							:      Union[str, Any]  						=							hf_model.state_dict()
     _lowercase							:      List[Any]  						=							{}

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

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

          if "weight" in torch_key and v.ndim == 2:
               _lowercase							:      Tuple  						=							v.T

          _lowercase							:      Union[str, Any]  						=							nn.Parameter(torch.from_numpy(lowerCAmelCase__))

     # Rename box prediction box flax params to pytorch HF
     _lowercase							:      Tuple  						=							flatten_nested_dict(flax_params['obj_box_head'])

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

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

          _lowercase							:      Optional[int]  						=							nn.Parameter(torch.from_numpy(lowerCAmelCase__))

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

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	str       ,  lowerCAmelCase__     :	List[Any])			->						Union[str, Any]:

     '''simple docstring'''




     for k, v in flax_attn_params.items():
          if k.startswith('transformer'):
               _lowercase							:      Tuple  						=							k.replace('transformer.resblocks'       ,  'text_model.encoder.layers')
          else:
               _lowercase							:      str  						=							k.replace('visual.transformer.resblocks'       ,  'vision_model.encoder.layers')

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

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

          if "weight" in torch_key and "out" in torch_key:
               _lowercase							:      str  						=							(v.shape[0] * v.shape[1], v.shape[2])
               _lowercase							:      str  						=							v.reshape(lowerCAmelCase__).T

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

          # Copy flax CLIP attn params to HF PyTorch params
          _lowercase							:      List[str]  						=							torch.from_numpy(lowerCAmelCase__)
          hf_backbone.state_dict()[torch_key].copy_(lowerCAmelCase__)

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	int)			->						Optional[int]:

     '''simple docstring'''




     _lowercase							:      Any  						=							{}
     _lowercase							:      int  						=							[]

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

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

def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[Any]       ,  lowerCAmelCase__     :	int)			->						Tuple:

     '''simple docstring'''




     _lowercase							:      Optional[int]  						=							CLIP(**lowerCAmelCase__)
     torch_model.eval()
     _lowercase							:      Any  						=							torch_model.state_dict()

     _lowercase							:      Union[str, Any]  						=							flatten_nested_dict(flax_params['backbone']['clip'])
     _lowercase							:      Union[str, Any]  						=							{}

     for flax_key, v in flax_clip_params.items():
          _lowercase							:      int  						=							flax_key.replace('/'       ,  '.')
          _lowercase							:      Union[str, 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							:      Union[str, Any]  						=							torch_key.replace('text_projection.kernel'       ,  'text_projection')
          _lowercase							:      Tuple  						=							torch_key.replace('visual.proj.kernel'       ,  'visual.proj')
          _lowercase							:      Union[str, Any]  						=							torch_key.replace('.scale'       ,  '.weight')
          _lowercase							:      int  						=							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							:      Dict  						=							v

     _lowercase							:      Optional[int]  						=							_convert_attn_layers(lowerCAmelCase__)
     new_torch_params.update(lowerCAmelCase__)
     _lowercase							:      Dict  						=							{}

     # 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_(lowerCAmelCase__)
          else:
               _lowercase							:      Optional[Any]  						=							param

     return torch_clip_params, torch_model, attn_params

@torch.no_grad()
def 		SCREAMING_SNAKE_CASE      (					lowerCAmelCase__     :	Optional[int]       ,  lowerCAmelCase__     :	Tuple       ,  lowerCAmelCase__     :	Union[str, Any]       ,  lowerCAmelCase__     :	Any       ,  lowerCAmelCase__     :	str=None)			->						Optional[int]:

     '''simple docstring'''




     _lowercase							:      List[str]  						=							Repository(lowerCAmelCase__       ,  clone_from=F'''google/{pytorch_dump_folder_path}''')
     repo.git_pull()

     if config_path is not None:
          _lowercase							:      Dict  						=							OwlViTConfig.from_pretrained(lowerCAmelCase__)
     else:
          _lowercase							:      int  						=							OwlViTConfig()

     _lowercase							:      Union[str, Any]  						=							OwlViTModel(lowerCAmelCase__).eval()
     _lowercase							:      List[str]  						=							OwlViTForObjectDetection(lowerCAmelCase__).eval()

     copy_text_model_and_projection(lowerCAmelCase__       ,  lowerCAmelCase__)
     copy_vision_model_and_projection(lowerCAmelCase__       ,  lowerCAmelCase__)
     _lowercase							:      Any  						=							pt_backbone.logit_scale
     copy_flax_attn_params(lowerCAmelCase__       ,  lowerCAmelCase__)

     _lowercase							:      str  						=							hf_backbone
     copy_class_merge_token(lowerCAmelCase__       ,  lowerCAmelCase__)
     copy_class_box_heads(lowerCAmelCase__       ,  lowerCAmelCase__)

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

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

     # Initialize processor
     _lowercase							:      Dict  						=							OwlViTProcessor(image_processor=lowerCAmelCase__       ,  tokenizer=lowerCAmelCase__)
     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__":
    A						    =							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.'''
    )
    A						    =							parser.parse_args()

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

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

    # Convert CLIP backbone
    A    ,       A    ,       A						    =							convert_clip_backbone(flax_params, torch_config)

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






'''simple docstring'''







import warnings

from ...utils import logging
from .image_processing_mobilenet_va import MobileNetVaImageProcessor


A						    =							logging.get_logger(__name__)



class 							__SCREAMING_SNAKE_CASE			(       lowerCAmelCase__      ):
      '''simple docstring'''



      def __init__(							self    : Tuple   ,*UpperCamelCase    : str   ,**UpperCamelCase    : List[Any]  )		->  None:
           warnings.warn(
               'The class MobileNetV2FeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
               ' Please use MobileNetV2ImageProcessor instead.'   ,UpperCamelCase   ,)
           super().__init__(*UpperCamelCase   ,**UpperCamelCase  )