


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      			=  {
    """vit_b32""": {
        """embed_dim""": 512,
        """image_resolution""": 768,
        """context_length""": 16,
        """vocab_size""": 4_9408,
        """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""": 4_9408,
        """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""": 4_9408,
        """vision_layers""": 24,
        """vision_width""": 1024,
        """vision_patch_size""": 14,
        """transformer_width""": 768,
        """transformer_heads""": 12,
        """transformer_layers""": 12,
    },
}
def 	_lowerCamelCase  (		A_  :	List[Any]							,  A_  :	Dict=""							,  A_  :	Dict="/"      )			->			Tuple:

   '''simple docstring'''






   UpperCamelCase__ :			Optional[Any] 						=[]

   for k, v in params.items():
      UpperCamelCase__ :			Tuple 						=parent_key + sep + k if parent_key else k

      if isinstance(a__							,  collections.MutableMapping      ):
         items.extend(flatten_nested_dict(a__							,  a__							,  sep=a__      ).items()      )
      else:
         items.append((new_key, v)      )
   return dict(a__      )
def 	_lowerCamelCase  (		A_  :	Any      )			->			List[str]:

   '''simple docstring'''






   return jax.tree_util.tree_map(lambda A_      : x.astype(jnp.floataa      ) if x.dtype == jnp.bfloataa else x							,  a__      )
def 	_lowerCamelCase  (		A_  :	Union[str, Any]							,  A_  :	Any      )			->			List[Any]:

   '''simple docstring'''






   UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__ :			int 						=pt_attn_layer.in_proj_weight.chunk(3							,  dim=0      )
   UpperCamelCase__						,	UpperCamelCase__						,	UpperCamelCase__ :			Any 						=pt_attn_layer.in_proj_bias.chunk(3							,  dim=0      )

   UpperCamelCase__ :			List[str] 						=pt_attn_layer.out_proj.weight
   UpperCamelCase__ :			Any 						=pt_attn_layer.out_proj.bias

   UpperCamelCase__ :			List[str] 						=q_proj
   UpperCamelCase__ :			Any 						=q_proj_bias

   UpperCamelCase__ :			int 						=k_proj
   UpperCamelCase__ :			Optional[Any] 						=k_proj_bias

   UpperCamelCase__ :			Tuple 						=v_proj
   UpperCamelCase__ :			Tuple 						=v_proj_bias

   UpperCamelCase__ :			Tuple 						=out_proj_weights
   UpperCamelCase__ :			List[str] 						=out_proj_bias
def 	_lowerCamelCase  (		A_  :	Dict							,  A_  :	List[Any]      )			->			Dict:

   '''simple docstring'''






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

   '''simple docstring'''






   UpperCamelCase__ :			List[str] 						=pt_linear.weight
   UpperCamelCase__ :			List[Any] 						=pt_linear.bias
def 	_lowerCamelCase  (		A_  :	List[str]							,  A_  :	Optional[Any]      )			->			List[str]:

   '''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 	_lowerCamelCase  (		A_  :	Optional[Any]							,  A_  :	int      )			->			List[str]:

   '''simple docstring'''






   for hf_layer, pt_layer in zip(a__							,  a__      ):
      copy_layer(a__							,  a__      )
def 	_lowerCamelCase  (		A_  :	List[Any]							,  A_  :	List[Any]      )			->			Dict:

   '''simple docstring'''






   UpperCamelCase__ :			Optional[Any] 						=pt_model.token_embedding.weight
   UpperCamelCase__ :			Optional[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 	_lowerCamelCase  (		A_  :	Tuple							,  A_  :	Dict      )			->			List[Any]:

   '''simple docstring'''






   UpperCamelCase__ :			Optional[Any] 						=pt_model.text_projection.data.T

   # copy text encoder
   copy_encoder(hf_model.text_model							,  a__      )
def 	_lowerCamelCase  (		A_  :	List[str]							,  A_  :	Optional[Any]      )			->			int:

   '''simple docstring'''






   UpperCamelCase__ :			Optional[int] 						=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__ :			Any 						=pt_model.visual.conva.weight.data
   UpperCamelCase__ :			int 						=pt_model.visual.class_embedding
   UpperCamelCase__ :			Optional[int] 						=pt_model.visual.positional_embedding.data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers							,  pt_model.visual.transformer.resblocks      )
def 	_lowerCamelCase  (		A_  :	Any							,  A_  :	int      )			->			List[str]:

   '''simple docstring'''






   UpperCamelCase__ :			Union[str, Any] 						=flatten_nested_dict(flax_params["backbone"]["merged_class_token"]      )

   UpperCamelCase__ :			Dict 						=torch.from_numpy(flax_class_token_params["scale"]      )
   UpperCamelCase__ :			int 						=torch.from_numpy(flax_class_token_params["bias"]      )
   UpperCamelCase__ :			Any 						=nn.Parameter(a__      )
   UpperCamelCase__ :			str 						=nn.Parameter(a__      )
def 	_lowerCamelCase  (		A_  :	Union[str, Any]							,  A_  :	Any      )			->			Dict:

   '''simple docstring'''






   UpperCamelCase__ :			str 						=hf_model.state_dict()
   UpperCamelCase__ :			Union[str, Any] 						={}

   # Rename class prediction head flax params to pytorch HF
   UpperCamelCase__ :			Union[str, Any] 						=flatten_nested_dict(flax_params["class_head"]      )

   for flax_key, v in flax_class_params.items():
      UpperCamelCase__ :			Optional[Any] 						=flax_key.replace("/"							,  "."      )
      UpperCamelCase__ :			Any 						=torch_key.replace(".kernel"							,  ".weight"      )
      UpperCamelCase__ :			Tuple 						=torch_key.replace("Dense_0"							,  "dense0"      )
      UpperCamelCase__ :			Optional[int] 						="class_head." + torch_key

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

      UpperCamelCase__ :			str 						=nn.Parameter(torch.from_numpy(a__      )      )

   # Rename box prediction box flax params to pytorch HF
   UpperCamelCase__ :			List[str] 						=flatten_nested_dict(flax_params["obj_box_head"]      )

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

      if "weight" in torch_key and v.ndim == 2:
         UpperCamelCase__ :			List[str] 						=v.T

      UpperCamelCase__ :			List[str] 						=nn.Parameter(torch.from_numpy(a__      )      )

   # Copy flax params to PyTorch params
   for name, param in new_params.items():
      if name in pt_params.keys():
         pt_params[name].copy_(a__      )
def 	_lowerCamelCase  (		A_  :	List[Any]							,  A_  :	List[str]      )			->			Tuple:

   '''simple docstring'''






   for k, v in flax_attn_params.items():
      if k.startswith("transformer"      ):
         UpperCamelCase__ :			List[Any] 						=k.replace("transformer.resblocks"							,  "text_model.encoder.layers"      )
      else:
         UpperCamelCase__ :			List[Any] 						=k.replace("visual.transformer.resblocks"							,  "vision_model.encoder.layers"      )

      UpperCamelCase__ :			Union[str, Any] 						=torch_key.replace("attn"							,  "self_attn"      )
      UpperCamelCase__ :			Union[str, Any] 						=torch_key.replace("key"							,  "k_proj"      )
      UpperCamelCase__ :			Union[str, Any] 						=torch_key.replace("value"							,  "v_proj"      )
      UpperCamelCase__ :			Optional[Any] 						=torch_key.replace("query"							,  "q_proj"      )
      UpperCamelCase__ :			Dict 						=torch_key.replace("out"							,  "out_proj"      )

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

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

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

      # Copy flax CLIP attn params to HF PyTorch params
      UpperCamelCase__ :			Dict 						=torch.from_numpy(a__      )
      hf_backbone.state_dict()[torch_key].copy_(a__      )
def 	_lowerCamelCase  (		A_  :	Dict      )			->			Optional[int]:

   '''simple docstring'''






   UpperCamelCase__ :			Dict 						={}
   UpperCamelCase__ :			Tuple 						=[]

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

         processed_attn_layers.append(a__      )
         UpperCamelCase__ :			Union[str, Any] 						=params[base + "out.weight"].shape[-1]
         UpperCamelCase__ :			Tuple 						=params[base + "out.weight"].reshape(a__							,  a__      ).T
         UpperCamelCase__ :			Union[str, Any] 						=params[base + "out.bias"]
      else:
         UpperCamelCase__ :			Union[str, Any] 						=v
   return new_params
def 	_lowerCamelCase  (		A_  :	Tuple							,  A_  :	Dict      )			->			int:

   '''simple docstring'''






   UpperCamelCase__ :			str 						=CLIP(**a__      )
   torch_model.eval()
   UpperCamelCase__ :			Optional[int] 						=torch_model.state_dict()

   UpperCamelCase__ :			Optional[Any] 						=flatten_nested_dict(flax_params["backbone"]["clip"]      )
   UpperCamelCase__ :			List[str] 						={}

   for flax_key, v in flax_clip_params.items():
      UpperCamelCase__ :			Any 						=flax_key.replace("/"							,  "."      )
      UpperCamelCase__ :			Dict 						=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__ :			Optional[Any] 						=torch_key[5:]

      UpperCamelCase__ :			Tuple 						=torch_key.replace("text_projection.kernel"							,  "text_projection"      )
      UpperCamelCase__ :			int 						=torch_key.replace("visual.proj.kernel"							,  "visual.proj"      )
      UpperCamelCase__ :			Any 						=torch_key.replace(".scale"							,  ".weight"      )
      UpperCamelCase__ :			str 						=torch_key.replace(".kernel"							,  ".weight"      )

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

      UpperCamelCase__ :			str 						=v

   UpperCamelCase__ :			Optional[int] 						=_convert_attn_layers(a__      )
   new_torch_params.update(a__      )
   UpperCamelCase__ :			Optional[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_(a__      )
      else:
         UpperCamelCase__ :			Dict 						=param

   return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def 	_lowerCamelCase  (		A_  :	List[str]							,  A_  :	Union[str, Any]							,  A_  :	Optional[int]							,  A_  :	int							,  A_  :	List[Any]=None      )			->			Tuple:

   '''simple docstring'''






   UpperCamelCase__ :			Union[str, Any] 						=Repository(a__							,  clone_from=f'''google/{pytorch_dump_folder_path}'''      )
   repo.git_pull()

   if config_path is not None:
      UpperCamelCase__ :			int 						=OwlViTConfig.from_pretrained(a__      )
   else:
      UpperCamelCase__ :			str 						=OwlViTConfig()

   UpperCamelCase__ :			Optional[int] 						=OwlViTModel(a__      ).eval()
   UpperCamelCase__ :			Any 						=OwlViTForObjectDetection(a__      ).eval()

   copy_text_model_and_projection(a__							,  a__      )
   copy_vision_model_and_projection(a__							,  a__      )
   UpperCamelCase__ :			Any 						=pt_backbone.logit_scale
   copy_flax_attn_params(a__							,  a__      )

   UpperCamelCase__ :			Dict 						=hf_backbone
   copy_class_merge_token(a__							,  a__      )
   copy_class_box_heads(a__							,  a__      )

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

   # Initialize image processor
   UpperCamelCase__ :			Tuple 						=OwlViTImageProcessor(
       size=config.vision_config.image_size							,  crop_size=config.vision_config.image_size      )
   # Initialize tokenizer
   UpperCamelCase__ :			List[str] 						=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32"							,  pad_token="!"							,  model_max_length=1_6      )

   # Initialize processor
   UpperCamelCase__ :			Tuple 						=OwlViTProcessor(image_processor=a__							,  tokenizer=a__      )
   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      			=  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      			=  parser.parse_args()

     # Initialize PyToch clip model
     __UpperCAmelCase      			=  args.owlvit_version
     if model_name == "clip_b16":
          __UpperCAmelCase      			=  CONFIGS["""vit_b16"""]
     elif model_name == "clip_b32":
          __UpperCAmelCase      			=  CONFIGS["""vit_b32"""]
     elif model_name == "clip_l14":
          __UpperCAmelCase      			=  CONFIGS["""vit_l14"""]

     # Load from checkpoint and convert params to float-32
     __UpperCAmelCase      			=  checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)["""optimizer"""]["""target"""]
     __UpperCAmelCase      			=  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      			=  convert_clip_backbone(flax_params, torch_config)

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







'''simple docstring'''
import numpy as np
def 							snake_case_    (  a__						:      Tuple     ):



    """simple docstring"""




    __lowercase							,__lowercase     							=		a.shape
    __lowercase     							=		min(a__       ,a__     )
    __lowercase     							=		np.eye(a__     )
    __lowercase     							=		a.copy()

    for k in range(t - 1     ):
        # select a column of modified matrix A':
        __lowercase     							=		r[k:, [k]]
        # construct first basis vector
        __lowercase     							=		np.zeros_like(a__     )
        __lowercase     							=		1.0
        # determine scaling factor
        __lowercase     							=		np.linalg.norm(a__     )
        # construct vector v for Householder reflection
        __lowercase     							=		x + np.sign(x[0]     ) * alpha * ea
        v /= np.linalg.norm(a__     )

        # construct the Householder matrix
        __lowercase     							=		np.eye(m - k     ) - 2.0 * v @ v.T
        # pad with ones and zeros as necessary
        __lowercase     							=		np.block([[np.eye(a__     ), np.zeros((k, m - k)     )], [np.zeros((m - k, k)     ), q_k]]     )

        __lowercase     							=		q @ q_k.T
        __lowercase     							=		q_k @ r

    return q, r


if __name__ == "__main__":
  import doctest

  doctest.testmod()




