



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_1_2,
        "image_resolution": 7_6_8,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 1_2,
        "vision_width": 7_6_8,
        "vision_patch_size": 3_2,
        "transformer_width": 5_1_2,
        "transformer_heads": 8,
        "transformer_layers": 1_2,
    },
    "vit_b16": {
        "embed_dim": 5_1_2,
        "image_resolution": 7_6_8,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 1_2,
        "vision_width": 7_6_8,
        "vision_patch_size": 1_6,
        "transformer_width": 5_1_2,
        "transformer_heads": 8,
        "transformer_layers": 1_2,
    },
    "vit_l14": {
        "embed_dim": 7_6_8,
        "image_resolution": 8_4_0,
        "context_length": 1_6,
        "vocab_size": 4_9_4_0_8,
        "vision_layers": 2_4,
        "vision_width": 1_0_2_4,
        "vision_patch_size": 1_4,
        "transformer_width": 7_6_8,
        "transformer_heads": 1_2,
        "transformer_layers": 1_2,
    },
}

def 	a						(				_UpperCAmelCase							: Optional[int]			,       _UpperCAmelCase							: List[str]=""			,       _UpperCAmelCase							: Union[str, Any]="/"						):







		'''simple docstring'''





		__UpperCAmelCase   :					str													=     []

		for k, v in params.items():
				__UpperCAmelCase   :					int													=     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 	a						(				_UpperCAmelCase							: Union[str, Any]						):







		'''simple docstring'''





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

def 	a						(				_UpperCAmelCase							: List[str]			,       _UpperCAmelCase							: str						):







		'''simple docstring'''





		__UpperCAmelCase							, __UpperCAmelCase							, __UpperCAmelCase   :					List[Any]													=     pt_attn_layer.in_proj_weight.chunk(3			,       dim=0						)
		__UpperCAmelCase							, __UpperCAmelCase							, __UpperCAmelCase   :					Optional[Any]													=     pt_attn_layer.in_proj_bias.chunk(3			,       dim=0						)

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

		__UpperCAmelCase   :					Union[str, Any]													=     q_proj
		__UpperCAmelCase   :					Any													=     q_proj_bias

		__UpperCAmelCase   :					int													=     k_proj
		__UpperCAmelCase   :					Union[str, Any]													=     k_proj_bias

		__UpperCAmelCase   :					Optional[Any]													=     v_proj
		__UpperCAmelCase   :					List[Any]													=     v_proj_bias

		__UpperCAmelCase   :					Union[str, Any]													=     out_proj_weights
		__UpperCAmelCase   :					int													=     out_proj_bias

def 	a						(				_UpperCAmelCase							: List[str]			,       _UpperCAmelCase							: str						):







		'''simple docstring'''





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

def 	a						(				_UpperCAmelCase							: List[Any]			,       _UpperCAmelCase							: int						):







		'''simple docstring'''





		__UpperCAmelCase   :					Dict													=     pt_linear.weight
		__UpperCAmelCase   :					Tuple													=     pt_linear.bias

def 	a						(				_UpperCAmelCase							: Union[str, Any]			,       _UpperCAmelCase							: int						):







		'''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 	a						(				_UpperCAmelCase							: Optional[Any]			,       _UpperCAmelCase							: Optional[int]						):







		'''simple docstring'''





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

def 	a						(				_UpperCAmelCase							: List[str]			,       _UpperCAmelCase							: List[str]						):







		'''simple docstring'''





		__UpperCAmelCase   :					List[Any]													=     pt_model.token_embedding.weight
		__UpperCAmelCase   :					List[str]													=     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 	a						(				_UpperCAmelCase							: Union[str, Any]			,       _UpperCAmelCase							: Optional[int]						):







		'''simple docstring'''





		__UpperCAmelCase   :					List[Any]													=     pt_model.text_projection.data.T

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

def 	a						(				_UpperCAmelCase							: Optional[Any]			,       _UpperCAmelCase							: Any						):







		'''simple docstring'''





		__UpperCAmelCase   :					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
		__UpperCAmelCase   :					Union[str, Any]													=     pt_model.visual.conva.weight.data
		__UpperCAmelCase   :					Optional[int]													=     pt_model.visual.class_embedding
		__UpperCAmelCase   :					Tuple													=     pt_model.visual.positional_embedding.data

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

def 	a						(				_UpperCAmelCase							: str			,       _UpperCAmelCase							: Dict						):







		'''simple docstring'''





		__UpperCAmelCase   :					Tuple													=     flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']						)

		__UpperCAmelCase   :					Any													=     torch.from_numpy(flax_class_token_params['''scale''']						)
		__UpperCAmelCase   :					List[Any]													=     torch.from_numpy(flax_class_token_params['''bias''']						)
		__UpperCAmelCase   :					Dict													=     nn.Parameter(_UpperCAmelCase						)
		__UpperCAmelCase   :					Dict													=     nn.Parameter(_UpperCAmelCase						)

def 	a						(				_UpperCAmelCase							: List[str]			,       _UpperCAmelCase							: Optional[Any]						):







		'''simple docstring'''





		__UpperCAmelCase   :					Optional[int]													=     hf_model.state_dict()
		__UpperCAmelCase   :					Dict													=     {}

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

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

				if "weight" in torch_key and v.ndim == 2:
						__UpperCAmelCase   :					Dict													=     v.T

				__UpperCAmelCase   :					int													=     nn.Parameter(torch.from_numpy(_UpperCAmelCase						)						)

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

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

				if "weight" in torch_key and v.ndim == 2:
						__UpperCAmelCase   :					str													=     v.T

				__UpperCAmelCase   :					Tuple													=     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 	a						(				_UpperCAmelCase							: List[Any]			,       _UpperCAmelCase							: Union[str, Any]						):







		'''simple docstring'''





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

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

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

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

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

				# Copy flax CLIP attn params to HF PyTorch params
				__UpperCAmelCase   :					Tuple													=     torch.from_numpy(_UpperCAmelCase						)
				hf_backbone.state_dict()[torch_key].copy_(_UpperCAmelCase						)

def 	a						(				_UpperCAmelCase							: Dict						):







		'''simple docstring'''





		__UpperCAmelCase   :					List[Any]													=     {}
		__UpperCAmelCase   :					int													=     []

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

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

def 	a						(				_UpperCAmelCase							: Dict			,       _UpperCAmelCase							: Union[str, Any]						):







		'''simple docstring'''





		__UpperCAmelCase   :					List[str]													=     CLIP(**_UpperCAmelCase						)
		torch_model.eval()
		__UpperCAmelCase   :					Tuple													=     torch_model.state_dict()

		__UpperCAmelCase   :					Optional[Any]													=     flatten_nested_dict(flax_params['''backbone''']['''clip''']						)
		__UpperCAmelCase   :					Dict													=     {}

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

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

				if "conv" in torch_key or "downsample.0.weight" in torch_key:
						__UpperCAmelCase   :					int													=     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   :					int													=     v

		__UpperCAmelCase   :					Optional[Any]													=     _convert_attn_layers(_UpperCAmelCase						)
		new_torch_params.update(_UpperCAmelCase						)
		__UpperCAmelCase   :					Union[str, Any]													=     {}

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

		return torch_clip_params, torch_model, attn_params

@torch.no_grad()
def 	a						(				_UpperCAmelCase							: Optional[Any]			,       _UpperCAmelCase							: Any			,       _UpperCAmelCase							: str			,       _UpperCAmelCase							: Tuple			,       _UpperCAmelCase							: Tuple=None						):







		'''simple docstring'''





		__UpperCAmelCase   :					Optional[Any]													=     Repository(_UpperCAmelCase			,       clone_from=f'google/{pytorch_dump_folder_path}'						)
		repo.git_pull()

		if config_path is not None:
				__UpperCAmelCase   :					Dict													=     OwlViTConfig.from_pretrained(_UpperCAmelCase						)
		else:
				__UpperCAmelCase   :					Tuple													=     OwlViTConfig()

		__UpperCAmelCase   :					Optional[int]													=     OwlViTModel(_UpperCAmelCase						).eval()
		__UpperCAmelCase   :					Any													=     OwlViTForObjectDetection(_UpperCAmelCase						).eval()

		copy_text_model_and_projection(_UpperCAmelCase			,       _UpperCAmelCase						)
		copy_vision_model_and_projection(_UpperCAmelCase			,       _UpperCAmelCase						)
		__UpperCAmelCase   :					Optional[int]													=     pt_backbone.logit_scale
		copy_flax_attn_params(_UpperCAmelCase			,       _UpperCAmelCase						)

		__UpperCAmelCase   :					Optional[int]													=     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
		__UpperCAmelCase   :					List[Any]													=     OwlViTImageProcessor(
		    size=config.vision_config.image_size			,       crop_size=config.vision_config.image_size						)
		# Initialize tokenizer
		__UpperCAmelCase   :					Tuple													=     CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''			,       pad_token='''!'''			,       model_max_length=16						)

		# Initialize processor
		__UpperCAmelCase   :					Optional[Any]													=     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__":
				__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)










import os
from shutil import copyfile
from typing import Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
				from .tokenization_reformer import ReformerTokenizer
else:
				__A             =None


__A             =logging.get_logger(__name__)


__A             ="▁"

__A             ={"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

__A             ={
    "vocab_file": {
        "google/reformer-crime-and-punishment": (
            "https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/spiece.model"
        )
    },
    "tokenizer_file": {
        "google/reformer-crime-and-punishment": (
            "https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/tokenizer.json"
        )
    },
}

__A             ={
    "google/reformer-crime-and-punishment": 5_2_4_2_8_8,
}

class 		UpperCAmelCase__      (      __UpperCamelCase       ):




			'''simple docstring'''






			UpperCamelCase      						=	VOCAB_FILES_NAMES
			UpperCamelCase      						=	PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase      						=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase      						=	["""input_ids""", """attention_mask"""]
			UpperCamelCase      						=	ReformerTokenizer




			def __init__(       self							:			Tuple	,	a_							:			Optional[int]=None	,	a_							:			int=None	,	a_							:			Dict="</s>"	,	a_							:			Dict="<unk>"	,	a_							:			List[Any]=[]	,	**a_							:			int	,	):
					'''simple docstring'''





					super().__init__(
					    a_	,	tokenizer_file=a_	,	eos_token=a_	,	unk_token=a_	,	additional_special_tokens=a_	,	**a_	,	)

					__UpperCAmelCase   :					Optional[Any]													=     vocab_file
					__UpperCAmelCase   :					Optional[Any]													=     False if not self.vocab_file else True




			def 			snake_case__					(       self							:			Dict	,	a_							:			str	,	a_							:			Optional[str] = None			):
					'''simple docstring'''





					if not self.can_save_slow_tokenizer:
							raise ValueError(
							    '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
							    '''tokenizer.'''			)

					if not os.path.isdir(a_			):
							logger.error(F'Vocabulary path ({save_directory}) should be a directory'			)
							return
					__UpperCAmelCase   :					Tuple													=     os.path.join(
					    a_	,	(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']			)

					if os.path.abspath(self.vocab_file			) != os.path.abspath(a_			):
							copyfile(self.vocab_file	,	a_			)

					return (out_vocab_file,)





