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


UpperCamelCase      :							List[str]          =					{
    '''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		(							snake_case :Optional[Any]							,   snake_case :List[str]=""							,   snake_case :List[Any]="/"						) ->      Tuple:
 __UpperCamelCase							    =			[]

 for k, v in params.items():
  __UpperCamelCase							    =			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 				A		(							snake_case :Any						) ->      Dict:
 return jax.tree_util.tree_map(lambda snake_case						: x.astype(jnp.floataa						) if x.dtype == jnp.bfloataa else x							,   snake_case						)






def 				A		(							snake_case :str							,   snake_case :List[str]						) ->      List[str]:
 __UpperCamelCase							    =			pt_attn_layer.in_proj_weight.chunk(3							,   dim=0						)
 __UpperCamelCase							    =			pt_attn_layer.in_proj_bias.chunk(3							,   dim=0						)

 __UpperCamelCase							    =			pt_attn_layer.out_proj.weight
 __UpperCamelCase							    =			pt_attn_layer.out_proj.bias

 __UpperCamelCase							    =			q_proj
 __UpperCamelCase							    =			q_proj_bias

 __UpperCamelCase							    =			k_proj
 __UpperCamelCase							    =			k_proj_bias

 __UpperCamelCase							    =			v_proj
 __UpperCamelCase							    =			v_proj_bias

 __UpperCamelCase							    =			out_proj_weights
 __UpperCamelCase							    =			out_proj_bias






def 				A		(							snake_case :str							,   snake_case :Union[str, Any]						) ->      Dict:
 copy_linear(hf_mlp.fca							,   pt_mlp.c_fc						)
 copy_linear(hf_mlp.fca							,   pt_mlp.c_proj						)






def 				A		(							snake_case :int							,   snake_case :Optional[int]						) ->      Dict:
 __UpperCamelCase							    =			pt_linear.weight
 __UpperCamelCase							    =			pt_linear.bias






def 				A		(							snake_case :Optional[Any]							,   snake_case :Optional[Any]						) ->      Dict:
 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		(							snake_case :str							,   snake_case :Union[str, Any]						) ->      Optional[int]:



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






def 				A		(							snake_case :List[Any]							,   snake_case :str						) ->      Any:
 __UpperCamelCase							    =			pt_model.token_embedding.weight
 __UpperCamelCase							    =			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		(							snake_case :str							,   snake_case :Any						) ->      Optional[int]:
 __UpperCamelCase							    =			pt_model.text_projection.data.T

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






def 				A		(							snake_case :int							,   snake_case :Dict						) ->      Any:
 __UpperCamelCase							    =			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							    =			pt_model.visual.conva.weight.data
 __UpperCamelCase							    =			pt_model.visual.class_embedding
 __UpperCamelCase							    =			pt_model.visual.positional_embedding.data

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






def 				A		(							snake_case :int							,   snake_case :Union[str, Any]						) ->      int:
 __UpperCamelCase							    =			flatten_nested_dict(flax_params['backbone']['merged_class_token']						)

 __UpperCamelCase							    =			torch.from_numpy(flax_class_token_params['scale']						)
 __UpperCamelCase							    =			torch.from_numpy(flax_class_token_params['bias']						)
 __UpperCamelCase							    =			nn.Parameter(snake_case						)
 __UpperCamelCase							    =			nn.Parameter(snake_case						)






def 				A		(							snake_case :str							,   snake_case :Tuple						) ->      Union[str, Any]:
 __UpperCamelCase							    =			hf_model.state_dict()
 __UpperCamelCase							    =			{}

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

 for flax_key, v in flax_class_params.items():
  __UpperCamelCase							    =			flax_key.replace('/'							,   '.'						)
  __UpperCamelCase							    =			torch_key.replace('.kernel'							,   '.weight'						)
  __UpperCamelCase							    =			torch_key.replace('Dense_0'							,   'dense0'						)
  __UpperCamelCase							    =			"""class_head.""" + torch_key

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

  __UpperCamelCase							    =			nn.Parameter(torch.from_numpy(snake_case						)						)

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

 for flax_key, v in flax_box_params.items():
  __UpperCamelCase							    =			flax_key.replace('/'							,   '.'						)
  __UpperCamelCase							    =			torch_key.replace('.kernel'							,   '.weight'						)
  __UpperCamelCase							    =			torch_key.replace('_'							,   ''						).lower()
  __UpperCamelCase							    =			"""box_head.""" + torch_key

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

  __UpperCamelCase							    =			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 				A		(							snake_case :Optional[Any]							,   snake_case :Optional[int]						) ->      List[str]:



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

  __UpperCamelCase							    =			torch_key.replace('attn'							,   'self_attn'						)
  __UpperCamelCase							    =			torch_key.replace('key'							,   'k_proj'						)
  __UpperCamelCase							    =			torch_key.replace('value'							,   'v_proj'						)
  __UpperCamelCase							    =			torch_key.replace('query'							,   'q_proj'						)
  __UpperCamelCase							    =			torch_key.replace('out'							,   'out_proj'						)

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

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

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

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






def 				A		(							snake_case :Optional[int]						) ->      Union[str, Any]:
 __UpperCamelCase							    =			{}
 __UpperCamelCase							    =			[]

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

   processed_attn_layers.append(snake_case						)
   __UpperCamelCase							    =			params[base + """out.weight"""].shape[-1]
   __UpperCamelCase							    =			params[base + """out.weight"""].reshape(snake_case							,   snake_case						).T
   __UpperCamelCase							    =			params[base + """out.bias"""]
  else:
   __UpperCamelCase							    =			v
 return new_params






def 				A		(							snake_case :int							,   snake_case :Tuple						) ->      Any:
 __UpperCamelCase							    =			CLIP(**snake_case						)
 torch_model.eval()
 __UpperCamelCase							    =			torch_model.state_dict()

 __UpperCamelCase							    =			flatten_nested_dict(flax_params['backbone']['clip']						)
 __UpperCamelCase							    =			{}

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

  __UpperCamelCase							    =			torch_key.replace('text_projection.kernel'							,   'text_projection'						)
  __UpperCamelCase							    =			torch_key.replace('visual.proj.kernel'							,   'visual.proj'						)
  __UpperCamelCase							    =			torch_key.replace('.scale'							,   '.weight'						)
  __UpperCamelCase							    =			torch_key.replace('.kernel'							,   '.weight'						)

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

  __UpperCamelCase							    =			v

 __UpperCamelCase							    =			_convert_attn_layers(snake_case						)
 new_torch_params.update(snake_case						)
 __UpperCamelCase							    =			{}

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

 return torch_clip_params, torch_model, attn_params






@torch.no_grad()
def 				A		(							snake_case :Tuple							,   snake_case :Any							,   snake_case :Union[str, Any]							,   snake_case :List[Any]							,   snake_case :Tuple=None						) ->      int:
 __UpperCamelCase							    =			Repository(snake_case							,   clone_from=f'google/{pytorch_dump_folder_path}'						)
 repo.git_pull()

 if config_path is not None:
  __UpperCamelCase							    =			OwlViTConfig.from_pretrained(snake_case						)
 else:
  __UpperCamelCase							    =			OwlViTConfig()

 __UpperCamelCase							    =			OwlViTModel(snake_case						).eval()
 __UpperCamelCase							    =			OwlViTForObjectDetection(snake_case						).eval()

 copy_text_model_and_projection(snake_case							,   snake_case						)
 copy_vision_model_and_projection(snake_case							,   snake_case						)
 __UpperCamelCase							    =			pt_backbone.logit_scale
 copy_flax_attn_params(snake_case							,   snake_case						)

 __UpperCamelCase							    =			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
 __UpperCamelCase							    =			OwlViTImageProcessor(
     size=config.vision_config.image_size							,   crop_size=config.vision_config.image_size						)
 # Initialize tokenizer
 __UpperCamelCase							    =			CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'							,   pad_token='!'							,   model_max_length=1_6						)

 # Initialize processor
 __UpperCamelCase							    =			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__":
      UpperCamelCase      :							Tuple          =					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      :							Union[str, Any]          =					parser.parse_args()

      # Initialize PyToch clip model
      UpperCamelCase      :							int          =					args.owlvit_version
      if model_name == "clip_b16":
            UpperCamelCase      :							Dict          =					CONFIGS['''vit_b16''']
      elif model_name == "clip_b32":
            UpperCamelCase      :							Any          =					CONFIGS['''vit_b32''']
      elif model_name == "clip_l14":
            UpperCamelCase      :							Union[str, Any]          =					CONFIGS['''vit_l14''']

      # Load from checkpoint and convert params to float-32
      UpperCamelCase      :							Union[str, Any]          =					checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['''optimizer''']['''target''']
      UpperCamelCase      :							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      :							Tuple          =					convert_clip_backbone(flax_params, torch_config)

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











# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

from .config_args import default_config_file, load_config_from_file
from .config_utils import SubcommandHelpFormatter


UpperCAmelCase_					:		int           = '''Update an existing config file with the latest defaults while maintaining the old configuration.'''

def   SCREAMING_SNAKE_CASE_				(				__magic_name__   :    str   )						->     Dict:






	"""simple docstring"""



	UpperCamelCase	:Optional[int]		=      args.config_file
	if config_file is None and Path(__magic_name__   ).exists():
		UpperCamelCase	:int		=      default_config_file
	elif not Path(__magic_name__   ).exists():
		raise ValueError(f"""The passed config file located at {config_file} doesn't exist."""   )
	UpperCamelCase	:Dict		=      load_config_from_file(__magic_name__   )

	if config_file.endswith(""".json"""   ):
		config.to_json_file(__magic_name__   )
	else:
		config.to_yaml_file(__magic_name__   )
	return config_file

def   SCREAMING_SNAKE_CASE_				(				__magic_name__   :    int      ,     __magic_name__   :    Optional[Any]   )						->     List[str]:






	"""simple docstring"""



	UpperCamelCase	:Union[str, Any]		=      parser.add_parser("""update"""      ,     parents=__magic_name__      ,     help=__magic_name__      ,     formatter_class=__magic_name__   )
	parser.add_argument(
	    """--config_file"""      ,     default=__magic_name__      ,     help=(
	        """The path to the config file to update. Will default to a file named default_config.yaml in the cache """
	        """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
	        """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
	        """with 'huggingface'."""
	    )      ,     )

	parser.set_defaults(func=__magic_name__   )
	return parser

def   SCREAMING_SNAKE_CASE_				(				__magic_name__   :    Union[str, Any]   )						->     Union[str, Any]:






	"""simple docstring"""



	UpperCamelCase	:Tuple		=      update_config(__magic_name__   )
	print(f"""Sucessfully updated the configuration file at {config_file}."""   )
