code: 
  
  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_       :				Optional[Any]		   =      {
      'vit_b32': {
          'embed_dim': 512,
          'image_resolution': 768,
          'context_length': 16,
          'vocab_size': 49_408,
          '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': 49_408,
          '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': 49_408,
          'vision_layers': 24,
          'vision_width': 1_024,
          'vision_patch_size': 14,
          'transformer_width': 768,
          'transformer_heads': 12,
          'transformer_layers': 12,
      },
  }
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[Any] ,a__						:      Dict="" ,a__						:      int="/"					)				->							Optional[Any]:
  	__A   : List[str]								= []
  
  	for k, v in params.items():
  		__A   : Tuple								= parent_key + sep + k if parent_key else k
  
  		if isinstance(_SCREAMING_SNAKE_CASE ,collections.MutableMapping					):
  			items.extend(flatten_nested_dict(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,sep=_SCREAMING_SNAKE_CASE					).items()					)
  		else:
  			items.append((new_key, v)					)
  	return dict(_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Tuple					)				->							List[Any]:
  	return jax.tree_util.tree_map(lambda a__					: x.astype(jnp.floataa					) if x.dtype == jnp.bfloataa else x ,_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Any ,a__						:      Optional[Any]					)				->							str:
  	__A   : str								= pt_attn_layer.in_proj_weight.chunk(3 ,dim=0					)
  	__A   : Any								= pt_attn_layer.in_proj_bias.chunk(3 ,dim=0					)
  
  	__A   : str								= pt_attn_layer.out_proj.weight
  	__A   : str								= pt_attn_layer.out_proj.bias
  
  	__A   : List[Any]								= q_proj
  	__A   : Dict								= q_proj_bias
  
  	__A   : Optional[int]								= k_proj
  	__A   : Tuple								= k_proj_bias
  
  	__A   : str								= v_proj
  	__A   : Tuple								= v_proj_bias
  
  	__A   : List[str]								= out_proj_weights
  	__A   : Optional[Any]								= out_proj_bias
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[int] ,a__						:      Union[str, Any]					)				->							Tuple:
  	copy_linear(hf_mlp.fca ,pt_mlp.c_fc					)
  	copy_linear(hf_mlp.fca ,pt_mlp.c_proj					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      List[str] ,a__						:      Optional[int]					)				->							Optional[Any]:
  	__A   : Dict								= pt_linear.weight
  	__A   : Optional[int]								= pt_linear.bias
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[Any] ,a__						:      List[str]					)				->							Optional[Any]:
  	# 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        __SCREAMING_SNAKE_CASE    (      a__						:      Any ,a__						:      Dict					)				->							Any:
  	for hf_layer, pt_layer in zip(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					):
  		copy_layer(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Dict ,a__						:      Dict					)				->							Any:
  	# copy  embeds
  	__A   : Optional[int]								= pt_model.token_embedding.weight
  	__A   : int								= 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    (      a__						:      str ,a__						:      Optional[Any]					)				->							Any:
  	# copy projection
  	__A   : List[Any]								= pt_model.text_projection.data.T
  
  	# copy text encoder
  	copy_encoder(hf_model.text_model ,_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      List[str] ,a__						:      Union[str, Any]					)				->							int:
  	# copy projection
  	__A   : 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
  	__A   : Optional[int]								= pt_model.visual.conva.weight.data
  	__A   : Union[str, Any]								= pt_model.visual.class_embedding
  	__A   : Union[str, 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    (      a__						:      List[str] ,a__						:      List[Any]					)				->							Union[str, Any]:
  	__A   : Dict								= flatten_nested_dict(flax_params["""backbone"""]["""merged_class_token"""]					)
  
  	__A   : Union[str, Any]								= torch.from_numpy(flax_class_token_params["""scale"""]					)
  	__A   : Any								= torch.from_numpy(flax_class_token_params["""bias"""]					)
  	__A   : int								= nn.Parameter(_SCREAMING_SNAKE_CASE					)
  	__A   : List[str]								= nn.Parameter(_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[int] ,a__						:      List[Any]					)				->							str:
  	__A   : int								= hf_model.state_dict()
  	__A   : Union[str, Any]								= {}
  
  	# Rename class prediction head flax params to pytorch HF
  	__A   : str								= flatten_nested_dict(flax_params["""class_head"""]					)
  
  	for flax_key, v in flax_class_params.items():
  		__A   : Union[str, Any]								= flax_key.replace("""/""" ,"""."""					)
  		__A   : str								= torch_key.replace(""".kernel""" ,""".weight"""					)
  		__A   : str								= torch_key.replace("""Dense_0""" ,"""dense0"""					)
  		__A   : Optional[Any]								= 'class_head.' + torch_key
  
  		if "weight" in torch_key and v.ndim == 2:
  			__A   : Union[str, Any]								= v.T
  
  		__A   : List[Any]								= nn.Parameter(torch.from_numpy(_SCREAMING_SNAKE_CASE					)					)
  
  	# Rename box prediction box flax params to pytorch HF
  	__A   : Union[str, Any]								= flatten_nested_dict(flax_params["""obj_box_head"""]					)
  
  	for flax_key, v in flax_box_params.items():
  		__A   : str								= flax_key.replace("""/""" ,"""."""					)
  		__A   : Dict								= torch_key.replace(""".kernel""" ,""".weight"""					)
  		__A   : Optional[Any]								= torch_key.replace("""_""" ,""""""					).lower()
  		__A   : Optional[int]								= 'box_head.' + torch_key
  
  		if "weight" in torch_key and v.ndim == 2:
  			__A   : Optional[int]								= v.T
  
  		__A   : List[str]								= nn.Parameter(torch.from_numpy(_SCREAMING_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_(_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Any ,a__						:      Dict					)				->							Dict:
  	for k, v in flax_attn_params.items():
  		if k.startswith("""transformer"""					):
  			__A   : int								= k.replace("""transformer.resblocks""" ,"""text_model.encoder.layers"""					)
  		else:
  			__A   : Any								= k.replace("""visual.transformer.resblocks""" ,"""vision_model.encoder.layers"""					)
  
  		__A   : Optional[Any]								= torch_key.replace("""attn""" ,"""self_attn"""					)
  		__A   : Optional[int]								= torch_key.replace("""key""" ,"""k_proj"""					)
  		__A   : int								= torch_key.replace("""value""" ,"""v_proj"""					)
  		__A   : Union[str, Any]								= torch_key.replace("""query""" ,"""q_proj"""					)
  		__A   : Tuple								= torch_key.replace("""out""" ,"""out_proj"""					)
  
  		if "bias" in torch_key and v.ndim == 2:
  			__A   : int								= v.shape[0] * v.shape[1]
  			__A   : Optional[int]								= v.reshape(_SCREAMING_SNAKE_CASE					)
  
  		if "weight" in torch_key and "out" in torch_key:
  			__A   : Any								= (v.shape[0] * v.shape[1], v.shape[2])
  			__A   : Any								= v.reshape(_SCREAMING_SNAKE_CASE					).T
  
  		if "weight" in torch_key and "out" not in torch_key:
  			__A   : Any								= (v.shape[0], v.shape[1] * v.shape[2])
  			__A   : Union[str, Any]								= v.reshape(_SCREAMING_SNAKE_CASE					).T
  
  		# Copy flax CLIP attn params to HF PyTorch params
  		__A   : int								= torch.from_numpy(_SCREAMING_SNAKE_CASE					)
  		hf_backbone.state_dict()[torch_key].copy_(_SCREAMING_SNAKE_CASE					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[Any]					)				->							Any:
  	__A   : List[Any]								= {}
  	__A   : Union[str, Any]								= []
  
  	for k, v in params.items():
  		if "attn." in k:
  			__A   : List[str]								= k[: k.rindex("""attn."""					) + 5]
  			if base in processed_attn_layers:
  				continue
  
  			processed_attn_layers.append(_SCREAMING_SNAKE_CASE					)
  			__A   : Union[str, Any]								= params[base + 'out.weight'].shape[-1]
  			__A   : Tuple								= params[base + 'out.weight'].reshape(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					).T
  			__A   : Optional[int]								= params[base + 'out.bias']
  		else:
  			__A   : Tuple								= v
  	return new_params
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Dict ,a__						:      Union[str, Any]					)				->							Optional[Any]:
  	__A   : Union[str, Any]								= CLIP(**_SCREAMING_SNAKE_CASE					)
  	torch_model.eval()
  	__A   : Optional[Any]								= torch_model.state_dict()
  
  	__A   : int								= flatten_nested_dict(flax_params["""backbone"""]["""clip"""]					)
  	__A   : Optional[int]								= {}
  
  	for flax_key, v in flax_clip_params.items():
  		__A   : str								= flax_key.replace("""/""" ,"""."""					)
  		__A   : int								= 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"""					)
  		):
  			__A   : Union[str, Any]								= torch_key[5:]
  
  		__A   : Optional[int]								= torch_key.replace("""text_projection.kernel""" ,"""text_projection"""					)
  		__A   : List[Any]								= torch_key.replace("""visual.proj.kernel""" ,"""visual.proj"""					)
  		__A   : int								= torch_key.replace(""".scale""" ,""".weight"""					)
  		__A   : List[Any]								= torch_key.replace(""".kernel""" ,""".weight"""					)
  
  		if "conv" in torch_key or "downsample.0.weight" in torch_key:
  			__A   : Tuple								= 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
  			__A   : Optional[int]								= v.T
  
  		__A   : List[str]								= v
  
  	__A   : Union[str, Any]								= _convert_attn_layers(_SCREAMING_SNAKE_CASE					)
  	new_torch_params.update(_SCREAMING_SNAKE_CASE					)
  	__A   : str								= {}
  
  	# Copy flax CLIP backbone params to PyTorch params
  	for name, param in new_torch_params.items():
  		if name in torch_clip_params.keys():
  			__A   : List[Any]								= torch.from_numpy(new_torch_params[name]					)
  			torch_clip_params[name].copy_(_SCREAMING_SNAKE_CASE					)
  		else:
  			__A   : List[Any]								= param
  
  	return torch_clip_params, torch_model, attn_params
  @torch.no_grad()
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Any ,a__						:      Optional[int] ,a__						:      Optional[int] ,a__						:      Tuple ,a__						:      List[Any]=None					)				->							Optional[int]:
  	__A   : int								= Repository(_SCREAMING_SNAKE_CASE ,clone_from=f"""google/{pytorch_dump_folder_path}"""					)
  	repo.git_pull()
  
  	if config_path is not None:
  		__A   : Optional[Any]								= OwlViTConfig.from_pretrained(_SCREAMING_SNAKE_CASE					)
  	else:
  		__A   : List[Any]								= OwlViTConfig()
  
  	__A   : Union[str, Any]								= OwlViTModel(_SCREAMING_SNAKE_CASE					).eval()
  	__A   : Dict								= OwlViTForObjectDetection(_SCREAMING_SNAKE_CASE					).eval()
  
  	copy_text_model_and_projection(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  	copy_vision_model_and_projection(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  	__A   : Optional[int]								= pt_backbone.logit_scale
  	copy_flax_attn_params(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  
  	__A   : Any								= hf_backbone
  	copy_class_merge_token(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  	copy_class_box_heads(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE					)
  
  	# Save HF model
  	hf_model.save_pretrained(repo.local_dir					)
  
  	# Initialize image processor
  	__A   : Any								= OwlViTImageProcessor(
  	    size=config.vision_config.image_size ,crop_size=config.vision_config.image_size					)
  	# Initialize tokenizer
  	__A   : Any								= CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32""" ,pad_token="""!""" ,model_max_length=16					)
  
  	# Initialize processor
  	__A   : Union[str, Any]								= OwlViTProcessor(image_processor=_SCREAMING_SNAKE_CASE ,tokenizer=_SCREAMING_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__":
  			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_       :				List[Any]		   =      parser.parse_args()
  
  			# Initialize PyToch clip model
  			UpperCAmelCase_       :				Tuple		   =      args.owlvit_version
  			if model_name == "clip_b16":
  						UpperCAmelCase_       :				Tuple		   =      CONFIGS['vit_b16']
  			elif model_name == "clip_b32":
  						UpperCAmelCase_       :				Optional[int]		   =      CONFIGS['vit_b32']
  			elif model_name == "clip_l14":
  						UpperCAmelCase_       :				List[Any]		   =      CONFIGS['vit_l14']
  
  			# Load from checkpoint and convert params to float-32
  			UpperCAmelCase_       :				Optional[Any]		   =      checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['optimizer']['target']
  			UpperCAmelCase_       :				List[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_       :				str		   =      convert_clip_backbone(flax_params, torch_config)
  
  			convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)
  
  
  
  
  
  
style_context: 
  
  import unittest
  from pathlib import Path
  
  from transformers import is_vision_available, load_tool
  from transformers.testing_utils import get_tests_dir
  
  from .test_tools_common import ToolTesterMixin
  
  
  if is_vision_available():
  			from PIL import Image
  
  
  
  
  
  
  class     lowerCamelCase_			(  unittest.TestCase	,			_lowercase  ):
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Dict       ):
  		__A   : str								= load_tool("""image-captioning"""       )
  		self.tool.setup()
  		__A   : Dict								= load_tool("""image-captioning"""				,			remote=__A       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Optional[int]       ):
  		__A   : List[Any]								= Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""       )       ) / """000000039769.png"""       )
  		__A   : Optional[Any]								= self.tool(__A       )
  		self.assertEqual(__A				,			"""two cats sleeping on a couch"""       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Dict       ):
  		__A   : Any								= Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""       )       ) / """000000039769.png"""       )
  		__A   : List[str]								= self.remote_tool(__A       )
  		self.assertEqual(__A				,			"""two cats sleeping on a couch"""       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Union[str, Any]       ):
  		__A   : Optional[Any]								= Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""       )       ) / """000000039769.png"""       )
  		__A   : str								= self.tool(image=__A       )
  		self.assertEqual(__A				,			"""two cats sleeping on a couch"""       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    str       ):
  		__A   : Dict								= Image.open(Path(get_tests_dir("""fixtures/tests_samples/COCO"""       )       ) / """000000039769.png"""       )
  		__A   : List[Any]								= self.remote_tool(image=__A       )
  		self.assertEqual(__A				,			"""two cats sleeping on a couch"""       )
  
  
  
  
  
  