




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


__SCREAMING_SNAKE_CASE      :				Optional[Any]				      ={
    """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      lowerCAmelCase_     (					lowerCamelCase:    List[str]       ,			lowerCamelCase:    List[str]=""       ,			lowerCamelCase:    List[Any]="/"  )   ->     Optional[Any]:


      '''simple docstring'''




      lowerCamelCase__ 					=			[]

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

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

def      lowerCAmelCase_     (					lowerCamelCase:    int  )   ->     Dict:


      '''simple docstring'''




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

def      lowerCAmelCase_     (					lowerCamelCase:    int       ,			lowerCamelCase:    List[Any]  )   ->     str:


      '''simple docstring'''




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

      lowerCamelCase__ 					=			pt_attn_layer.out_proj.weight
      lowerCamelCase__ 					=			pt_attn_layer.out_proj.bias

      lowerCamelCase__ 					=			q_proj
      lowerCamelCase__ 					=			q_proj_bias

      lowerCamelCase__ 					=			k_proj
      lowerCamelCase__ 					=			k_proj_bias

      lowerCamelCase__ 					=			v_proj
      lowerCamelCase__ 					=			v_proj_bias

      lowerCamelCase__ 					=			out_proj_weights
      lowerCamelCase__ 					=			out_proj_bias

def      lowerCAmelCase_     (					lowerCamelCase:    Any       ,			lowerCamelCase:    List[str]  )   ->     List[str]:


      '''simple docstring'''




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

def      lowerCAmelCase_     (					lowerCamelCase:    int       ,			lowerCamelCase:    str  )   ->     Optional[int]:


      '''simple docstring'''




      lowerCamelCase__ 					=			pt_linear.weight
      lowerCamelCase__ 					=			pt_linear.bias

def      lowerCAmelCase_     (					lowerCamelCase:    List[Any]       ,			lowerCamelCase:    Optional[int]  )   ->     Union[str, Any]:


      '''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_     (					lowerCamelCase:    Dict       ,			lowerCamelCase:    Any  )   ->     List[str]:


      '''simple docstring'''




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

def      lowerCAmelCase_     (					lowerCamelCase:    str       ,			lowerCamelCase:    Optional[Any]  )   ->     Any:


      '''simple docstring'''




      lowerCamelCase__ 					=			pt_model.token_embedding.weight
      lowerCamelCase__ 					=			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_     (					lowerCamelCase:    List[Any]       ,			lowerCamelCase:    Optional[int]  )   ->     Any:


      '''simple docstring'''




      lowerCamelCase__ 					=			pt_model.text_projection.data.T

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

def      lowerCAmelCase_     (					lowerCamelCase:    str       ,			lowerCamelCase:    Union[str, Any]  )   ->     List[Any]:


      '''simple docstring'''




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

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

def      lowerCAmelCase_     (					lowerCamelCase:    List[Any]       ,			lowerCamelCase:    Optional[Any]  )   ->     Dict:


      '''simple docstring'''




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

      lowerCamelCase__ 					=			torch.from_numpy(flax_class_token_params['scale']  )
      lowerCamelCase__ 					=			torch.from_numpy(flax_class_token_params['bias']  )
      lowerCamelCase__ 					=			nn.Parameter(lowerCamelCase  )
      lowerCamelCase__ 					=			nn.Parameter(lowerCamelCase  )

def      lowerCAmelCase_     (					lowerCamelCase:    List[str]       ,			lowerCamelCase:    Optional[int]  )   ->     Optional[int]:


      '''simple docstring'''




      lowerCamelCase__ 					=			hf_model.state_dict()
      lowerCamelCase__ 					=			{}

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

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

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

            lowerCamelCase__ 					=			nn.Parameter(torch.from_numpy(lowerCamelCase  )  )

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

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

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

            lowerCamelCase__ 					=			nn.Parameter(torch.from_numpy(lowerCamelCase  )  )

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

def      lowerCAmelCase_     (					lowerCamelCase:    Any       ,			lowerCamelCase:    Dict  )   ->     str:


      '''simple docstring'''




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

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

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

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

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

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

def      lowerCAmelCase_     (					lowerCamelCase:    str  )   ->     Optional[Any]:


      '''simple docstring'''




      lowerCamelCase__ 					=			{}
      lowerCamelCase__ 					=			[]

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

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

def      lowerCAmelCase_     (					lowerCamelCase:    Tuple       ,			lowerCamelCase:    Any  )   ->     Dict:


      '''simple docstring'''




      lowerCamelCase__ 					=			CLIP(**lowerCamelCase  )
      torch_model.eval()
      lowerCamelCase__ 					=			torch_model.state_dict()

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

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

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

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

            lowerCamelCase__ 					=			v

      lowerCamelCase__ 					=			_convert_attn_layers(lowerCamelCase  )
      new_torch_params.update(lowerCamelCase  )
      lowerCamelCase__ 					=			{}

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

      return torch_clip_params, torch_model, attn_params

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


      '''simple docstring'''




      lowerCamelCase__ 					=			Repository(lowerCamelCase       ,			clone_from=f"""google/{pytorch_dump_folder_path}"""  )
      repo.git_pull()

      if config_path is not None:
            lowerCamelCase__ 					=			OwlViTConfig.from_pretrained(lowerCamelCase  )
      else:
            lowerCamelCase__ 					=			OwlViTConfig()

      lowerCamelCase__ 					=			OwlViTModel(lowerCamelCase  ).eval()
      lowerCamelCase__ 					=			OwlViTForObjectDetection(lowerCamelCase  ).eval()

      copy_text_model_and_projection(lowerCamelCase       ,			lowerCamelCase  )
      copy_vision_model_and_projection(lowerCamelCase       ,			lowerCamelCase  )
      lowerCamelCase__ 					=			pt_backbone.logit_scale
      copy_flax_attn_params(lowerCamelCase       ,			lowerCamelCase  )

      lowerCamelCase__ 					=			hf_backbone
      copy_class_merge_token(lowerCamelCase       ,			lowerCamelCase  )
      copy_class_box_heads(lowerCamelCase       ,			lowerCamelCase  )

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

      # Initialize image processor
      lowerCamelCase__ 					=			OwlViTImageProcessor(
          size=config.vision_config.image_size       ,			crop_size=config.vision_config.image_size  )
      # Initialize tokenizer
      lowerCamelCase__ 					=			CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'       ,			pad_token='!'       ,			model_max_length=16  )

      # Initialize processor
      lowerCamelCase__ 					=			OwlViTProcessor(image_processor=lowerCamelCase       ,			tokenizer=lowerCamelCase  )
      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__":
      __SCREAMING_SNAKE_CASE      :				Optional[Any]				      =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."""
      )
      __SCREAMING_SNAKE_CASE      :				Optional[int]				      =parser.parse_args()

      # Initialize PyToch clip model
      __SCREAMING_SNAKE_CASE      :				str				      =args.owlvit_version
      if model_name == "clip_b16":
            __SCREAMING_SNAKE_CASE      :				Optional[int]				      =CONFIGS["""vit_b16"""]
      elif model_name == "clip_b32":
            __SCREAMING_SNAKE_CASE      :				Optional[Any]				      =CONFIGS["""vit_b32"""]
      elif model_name == "clip_l14":
            __SCREAMING_SNAKE_CASE      :				List[str]				      =CONFIGS["""vit_l14"""]

      # Load from checkpoint and convert params to float-32
      __SCREAMING_SNAKE_CASE      :				Optional[int]				      =checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)["""optimizer"""]["""target"""]
      __SCREAMING_SNAKE_CASE      :				Tuple				      =jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
      del variables

      # Convert CLIP backbone
      __SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      ,					__SCREAMING_SNAKE_CASE      :				Optional[int]				      =convert_clip_backbone(flax_params, torch_config)

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





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
      import tensorflow as tf






class      A     (						TensorFormatter[Mapping, "tf.Tensor", Mapping]						):
     """simple docstring"""

     def __init__(      self					:				int		,					A__					:				Tuple=None		,					**A__					:				Tuple     )			->    Dict:
           super().__init__(features=A__     )
           lowerCamelCase__ 					=			tf_tensor_kwargs
           import tensorflow as tf  # noqa: import tf at initialization




     def 	__snake_case				(      self					:				List[Any]		,					A__					:				List[str]     )			->    Dict:
           import tensorflow as tf

           if isinstance(A__		,					A__     ) and column:
                 if all(
                     isinstance(A__		,					tf.Tensor     ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column     ):
                       return tf.stack(A__     )
                 elif all(
                     isinstance(A__		,					(tf.Tensor, tf.RaggedTensor)     ) and x.ndim == 1 and x.dtype == column[0].dtype
                     for x in column     ):
                       # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
                       return tf.ragged.stack(A__     )

           return column




     def 	__snake_case				(      self					:				Tuple		,					A__					:				Optional[int]     )			->    int:
           import tensorflow as tf

           if value is None:
                 return value

           lowerCamelCase__ 					=			{}

           if isinstance(A__		,					(np.number, np.ndarray)     ) and np.issubdtype(value.dtype		,					np.integer     ):
                 lowerCamelCase__ 					=			{'dtype': tf.intaa}
           elif isinstance(A__		,					(np.number, np.ndarray)     ) and np.issubdtype(value.dtype		,					np.floating     ):
                 lowerCamelCase__ 					=			{'dtype': tf.floataa}
           elif config.PIL_AVAILABLE and "PIL" in sys.modules:
                 import PIL.Image

                 if isinstance(A__		,					PIL.Image.Image     ):
                       lowerCamelCase__ 					=			np.asarray(A__     )

           return tf.convert_to_tensor(A__		,					**{**default_dtype, **self.tf_tensor_kwargs}     )




     def 	__snake_case				(      self					:				Optional[int]		,					A__					:				Dict     )			->    List[str]:
           import tensorflow as tf

           # support for torch, tf, jax etc.
           if config.TORCH_AVAILABLE and "torch" in sys.modules:
                 import torch

                 if isinstance(A__		,					torch.Tensor     ):
                       return self._tensorize(data_struct.detach().cpu().numpy()[()]     )
           if hasattr(A__		,					'__array__'     ) and not isinstance(A__		,					tf.Tensor     ):
                 lowerCamelCase__ 					=			data_struct.__array__()
           # support for nested types like struct of list of struct
           if isinstance(A__		,					np.ndarray     ):
                 if data_struct.dtype == object:  # tf tensors cannot be instantied from an array of objects
                       return self._consolidate([self.recursive_tensorize(A__     ) for substruct in data_struct]     )
           elif isinstance(A__		,					(list, tuple)     ):
                 return self._consolidate([self.recursive_tensorize(A__     ) for substruct in data_struct]     )
           return self._tensorize(A__     )




     def 	__snake_case				(      self					:				List[Any]		,					A__					:				dict     )			->    List[str]:
           return map_nested(self._recursive_tensorize		,					A__		,					map_list=A__     )




     def 	__snake_case				(      self					:				Union[str, Any]		,					A__					:				pa.Table     )			->    Mapping:
           lowerCamelCase__ 					=			self.numpy_arrow_extractor().extract_row(A__     )
           lowerCamelCase__ 					=			self.python_features_decoder.decode_row(A__     )
           return self.recursive_tensorize(A__     )




     def 	__snake_case				(      self					:				Optional[Any]		,					A__					:				pa.Table     )			->    "tf.Tensor":
           lowerCamelCase__ 					=			self.numpy_arrow_extractor().extract_column(A__     )
           lowerCamelCase__ 					=			self.python_features_decoder.decode_column(A__		,					pa_table.column_names[0]     )
           lowerCamelCase__ 					=			self.recursive_tensorize(A__     )
           lowerCamelCase__ 					=			self._consolidate(A__     )
           return column




     def 	__snake_case				(      self					:				int		,					A__					:				pa.Table     )			->    Mapping:
           lowerCamelCase__ 					=			self.numpy_arrow_extractor().extract_batch(A__     )
           lowerCamelCase__ 					=			self.python_features_decoder.decode_batch(A__     )
           lowerCamelCase__ 					=			self.recursive_tensorize(A__     )
           for column_name in batch:
                 lowerCamelCase__ 					=			self._consolidate(batch[column_name]     )
           return batch