






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': 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    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase=""    ,	UpperCAmelCase="/"     )     ->		Optional[int]:







							"""simple docstring"""





							_a  :    List[Any]			=						[]

							for k, v in params.items():
														_a  :    Optional[Any]			=						parent_key + sep + k if parent_key else k

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






def    UpperCamelCase__  (   UpperCAmelCase     )     ->		str:







							"""simple docstring"""





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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Optional[int]:







							"""simple docstring"""





							_a		,					_a		,					_a  :    Any			=						pt_attn_layer.in_proj_weight.chunk(3    ,	dim=0     )
							_a		,					_a		,					_a  :    Dict			=						pt_attn_layer.in_proj_bias.chunk(3    ,	dim=0     )

							_a  :    Dict			=						pt_attn_layer.out_proj.weight
							_a  :    Optional[Any]			=						pt_attn_layer.out_proj.bias

							_a  :    Any			=						q_proj
							_a  :    List[Any]			=						q_proj_bias

							_a  :    Any			=						k_proj
							_a  :    Dict			=						k_proj_bias

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

							_a  :    Optional[Any]			=						out_proj_weights
							_a  :    Optional[int]			=						out_proj_bias






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Optional[int]:







							"""simple docstring"""





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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		int:







							"""simple docstring"""





							_a  :    Union[str, Any]			=						pt_linear.weight
							_a  :    List[Any]			=						pt_linear.bias






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Dict:







							"""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    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Tuple:







							"""simple docstring"""





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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		List[Any]:







							"""simple docstring"""





							_a  :    List[str]			=						pt_model.token_embedding.weight
							_a  :    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    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Union[str, Any]:







							"""simple docstring"""





							_a  :    Any			=						pt_model.text_projection.data.T

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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		int:







							"""simple docstring"""





							_a  :    List[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  :    List[Any]			=						pt_model.visual.conva.weight.data
							_a  :    int			=						pt_model.visual.class_embedding
							_a  :    str			=						pt_model.visual.positional_embedding.data

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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		List[str]:







							"""simple docstring"""





							_a  :    List[str]			=						flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']     )

							_a  :    Tuple			=						torch.from_numpy(flax_class_token_params['''scale''']     )
							_a  :    Tuple			=						torch.from_numpy(flax_class_token_params['''bias''']     )
							_a  :    Optional[int]			=						nn.Parameter(UpperCAmelCase     )
							_a  :    Union[str, Any]			=						nn.Parameter(UpperCAmelCase     )






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		List[Any]:







							"""simple docstring"""





							_a  :    int			=						hf_model.state_dict()
							_a  :    Any			=						{}

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

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

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

														_a  :    int			=						nn.Parameter(torch.from_numpy(UpperCAmelCase     )     )

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

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

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

														_a  :    str			=						nn.Parameter(torch.from_numpy(UpperCAmelCase     )     )

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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		Union[str, Any]:







							"""simple docstring"""





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

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

														if "bias" in torch_key and v.ndim == 2:
																					_a  :    Tuple			=						v.shape[0] * v.shape[1]
																					_a  :    Tuple			=						v.reshape(UpperCAmelCase     )

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

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

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






def    UpperCamelCase__  (   UpperCAmelCase     )     ->		Tuple:







							"""simple docstring"""





							_a  :    Dict			=						{}
							_a  :    Dict			=						[]

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

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






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase     )     ->		List[Any]:







							"""simple docstring"""





							_a  :    Dict			=						CLIP(**UpperCAmelCase     )
							torch_model.eval()
							_a  :    List[str]			=						torch_model.state_dict()

							_a  :    List[str]			=						flatten_nested_dict(flax_params['''backbone''']['''clip''']     )
							_a  :    Any			=						{}

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

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

														if "conv" in torch_key or "downsample.0.weight" in torch_key:
																					_a  :    Union[str, 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
																					_a  :    int			=						v.T

														_a  :    Any			=						v

							_a  :    Dict			=						_convert_attn_layers(UpperCAmelCase     )
							new_torch_params.update(UpperCAmelCase     )
							_a  :    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():
																					_a  :    Optional[Any]			=						torch.from_numpy(new_torch_params[name]     )
																					torch_clip_params[name].copy_(UpperCAmelCase     )
														else:
																					_a  :    Union[str, Any]			=						param

							return torch_clip_params, torch_model, attn_params






@torch.no_grad()
def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase=None     )     ->		Union[str, Any]:







							"""simple docstring"""





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

							if config_path is not None:
														_a  :    str			=						OwlViTConfig.from_pretrained(UpperCAmelCase     )
							else:
														_a  :    Optional[Any]			=						OwlViTConfig()

							_a  :    Tuple			=						OwlViTModel(UpperCAmelCase     ).eval()
							_a  :    Optional[int]			=						OwlViTForObjectDetection(UpperCAmelCase     ).eval()

							copy_text_model_and_projection(UpperCAmelCase    ,	UpperCAmelCase     )
							copy_vision_model_and_projection(UpperCAmelCase    ,	UpperCAmelCase     )
							_a  :    Any			=						pt_backbone.logit_scale
							copy_flax_attn_params(UpperCAmelCase    ,	UpperCAmelCase     )

							_a  :    str			=						hf_backbone
							copy_class_merge_token(UpperCAmelCase    ,	UpperCAmelCase     )
							copy_class_box_heads(UpperCAmelCase    ,	UpperCAmelCase     )

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

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

							# Initialize processor
							_a  :    Dict			=						OwlViTProcessor(image_processor=UpperCAmelCase    ,	tokenizer=UpperCAmelCase     )
							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)







def    UpperCamelCase__  (   UpperCAmelCase     )     ->		tuple[int, int]:







							"""simple docstring"""





							if len(UpperCAmelCase     ) == 1:
														return a[0], a[0]
							else:
														# split A into half.
														_a  :    Union[str, Any]			=						a[: len(UpperCAmelCase     ) // 2]
														_a  :    Optional[Any]			=						a[len(UpperCAmelCase     ) // 2 :]

														# 2 sub problems, 1/2 of original size.
														_a		,					_a  :    Union[str, Any]			=						max_difference(UpperCAmelCase     )
														_a		,					_a  :    List[Any]			=						max_difference(UpperCAmelCase     )

														# get min of first and max of second
														# linear time
														_a  :    Union[str, Any]			=						min(UpperCAmelCase     )
														_a  :    Any			=						max(UpperCAmelCase     )

														# 3 cases, either (small1, big1),
														# (min_first, max_second), (small2, big2)
														# constant comparisons
														if biga - smalla > max_second - min_first and biga - smalla > biga - smalla:
																					return smalla, biga
														elif biga - smalla > max_second - min_first:
																					return smalla, biga
														else:
																					return min_first, max_second


if __name__ == "__main__":
						import doctest

						doctest.testmod()