





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   :							int     =	{
    """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 				__UpperCamelCase  (		snake_case    ,	snake_case=""    ,	snake_case="/"     )  ->						List[Any]:


   '''simple docstring'''







   __A				 =       []

   for k, v in params.items():
      __A				 =       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 				__UpperCamelCase  (		snake_case     )  ->						Dict:


   '''simple docstring'''







   return jax.tree_util.tree_map(lambda snake_case     : x.astype(jnp.floataa     ) if x.dtype == jnp.bfloataa else x    ,	snake_case     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Optional[int]:


   '''simple docstring'''







   __A ,						__A ,						__A				 =       pt_attn_layer.in_proj_weight.chunk(3    ,	dim=0     )
   __A ,						__A ,						__A				 =       pt_attn_layer.in_proj_bias.chunk(3    ,	dim=0     )

   __A				 =       pt_attn_layer.out_proj.weight
   __A				 =       pt_attn_layer.out_proj.bias

   __A				 =       q_proj
   __A				 =       q_proj_bias

   __A				 =       k_proj
   __A				 =       k_proj_bias

   __A				 =       v_proj
   __A				 =       v_proj_bias

   __A				 =       out_proj_weights
   __A				 =       out_proj_bias
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Any:


   '''simple docstring'''







   copy_linear(hf_mlp.fca    ,	pt_mlp.c_fc     )
   copy_linear(hf_mlp.fca    ,	pt_mlp.c_proj     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						List[str]:


   '''simple docstring'''







   __A				 =       pt_linear.weight
   __A				 =       pt_linear.bias
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						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  (		snake_case    ,	snake_case     )  ->						Optional[Any]:


   '''simple docstring'''







   for hf_layer, pt_layer in zip(snake_case    ,	snake_case     ):
      copy_layer(snake_case    ,	snake_case     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Tuple:


   '''simple docstring'''







   __A				 =       pt_model.token_embedding.weight
   __A				 =       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  (		snake_case    ,	snake_case     )  ->						List[str]:


   '''simple docstring'''







   __A				 =       pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model    ,	snake_case     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Optional[int]:


   '''simple docstring'''







   __A				 =       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				 =       pt_model.visual.conva.weight.data
   __A				 =       pt_model.visual.class_embedding
   __A				 =       pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers    ,	pt_model.visual.transformer.resblocks     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Any:


   '''simple docstring'''







   __A				 =       flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']     )

   __A				 =       torch.from_numpy(flax_class_token_params['''scale''']     )
   __A				 =       torch.from_numpy(flax_class_token_params['''bias''']     )
   __A				 =       nn.Parameter(snake_case     )
   __A				 =       nn.Parameter(snake_case     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Tuple:


   '''simple docstring'''







   __A				 =       hf_model.state_dict()
   __A				 =       {}

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

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

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

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

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

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

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

      __A				 =       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 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						Optional[int]:


   '''simple docstring'''







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

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

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

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

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

      # Copy flax CLIP attn params to HF PyTorch params
      __A				 =       torch.from_numpy(snake_case     )
      hf_backbone.state_dict()[torch_key].copy_(snake_case     )
def 				__UpperCamelCase  (		snake_case     )  ->						List[str]:


   '''simple docstring'''







   __A				 =       {}
   __A				 =       []

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

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


   '''simple docstring'''







   __A				 =       CLIP(**snake_case     )
   torch_model.eval()
   __A				 =       torch_model.state_dict()

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

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

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

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

      __A				 =       v

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

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

   return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def 				__UpperCamelCase  (		snake_case    ,	snake_case    ,	snake_case    ,	snake_case    ,	snake_case=None     )  ->						List[Any]:


   '''simple docstring'''







   __A				 =       Repository(snake_case    ,	clone_from=F"google/{pytorch_dump_folder_path}"     )
   repo.git_pull()

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

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

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

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

   # Initialize processor
   __A				 =       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   :							int     =	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   :							Optional[int]     =	parser.parse_args()

    # Initialize PyToch clip model
    _UpperCamelCase   :							Optional[Any]     =	args.owlvit_version
    if model_name == "clip_b16":
        _UpperCamelCase   :							Any     =	CONFIGS["""vit_b16"""]
    elif model_name == "clip_b32":
        _UpperCamelCase   :							Tuple     =	CONFIGS["""vit_b32"""]
    elif model_name == "clip_l14":
        _UpperCamelCase   :							Optional[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   :							List[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)








from __future__ import annotations

import sys
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						tuple[int, int]:


   '''simple docstring'''







   if abs(snake_case     ) == 1:
      return a, 0
   elif abs(snake_case     ) == 1:
      return 0, b

   __A ,						__A				 =       a, b
   __A ,						__A				 =       1, 0
   __A ,						__A				 =       0, 1

   while remainder != 0:
      __A				 =       old_remainder // remainder
      __A ,						__A				 =       remainder, old_remainder - quotient * remainder
      __A ,						__A				 =       coeff_a, old_coeff_a - quotient * coeff_a
      __A ,						__A				 =       coeff_b, old_coeff_b - quotient * coeff_b

   # sign correction for negative numbers
   if a < 0:
      __A				 =       -old_coeff_a
   if b < 0:
      __A				 =       -old_coeff_b

   return old_coeff_a, old_coeff_b
def 				__UpperCamelCase  (		)  ->						Tuple:


   '''simple docstring'''







   if len(sys.argv     ) < 3:
      print('''2 integer arguments required'''     )
      return 1
   __A				 =       int(sys.argv[1]     )
   __A				 =       int(sys.argv[2]     )
   print(extended_euclidean_algorithm(snake_case    ,	snake_case     )     )
   return 0


if __name__ == "__main__":
    raise SystemExit(main())

