





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


__lowercase													=			{
    """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 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Optional[int]     ,							__lowercase:		Tuple=""     ,							__lowercase:		Dict="/"):



  '''simple docstring'''






  __lowerCAmelCase			=				[]

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

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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Optional[Any]):



  '''simple docstring'''






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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Union[str, Any]     ,							__lowercase:		Optional[Any]):



  '''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 					__SCREAMING_SNAKE_CASE  (							__lowercase:		str     ,							__lowercase:		str):



  '''simple docstring'''






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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		int     ,							__lowercase:		List[Any]):



  '''simple docstring'''






  __lowerCAmelCase			=				pt_linear.weight
  __lowerCAmelCase			=				pt_linear.bias


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Union[str, Any]     ,							__lowercase:		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 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Tuple     ,							__lowercase:		Tuple):



  '''simple docstring'''






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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		List[Any]     ,							__lowercase:		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 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Any     ,							__lowercase:		Union[str, Any]):



  '''simple docstring'''






  __lowerCAmelCase			=				pt_model.text_projection.data.T

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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Optional[Any]     ,							__lowercase:		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 					__SCREAMING_SNAKE_CASE  (							__lowercase:		List[Any]     ,							__lowercase:		str):



  '''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(__lowercase)
  __lowerCAmelCase			=				nn.Parameter(__lowercase)


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		List[str]     ,							__lowercase:		Tuple):



  '''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(__lowercase))

  # 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(__lowercase))

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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		Union[str, Any]     ,							__lowercase:		Any):



  '''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(__lowercase)

    if "weight" in torch_key and "out" in torch_key:
      __lowerCAmelCase			=				(v.shape[0] * v.shape[1], v.shape[2])
      __lowerCAmelCase			=				v.reshape(__lowercase).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(__lowercase).T

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


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		List[str]):



  '''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(__lowercase)
      __lowerCAmelCase			=				params[base + '''out.weight'''].shape[-1]
      __lowerCAmelCase			=				params[base + '''out.weight'''].reshape(__lowercase     ,							__lowercase).T
      __lowerCAmelCase			=				params[base + '''out.bias''']
    else:
      __lowerCAmelCase			=				v
  return new_params


def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		List[str]     ,							__lowercase:		List[Any]):



  '''simple docstring'''






  __lowerCAmelCase			=				CLIP(**__lowercase)
  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(__lowercase)
  new_torch_params.update(__lowercase)
  __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_(__lowercase)
    else:
      __lowerCAmelCase			=				param

  return torch_clip_params, torch_model, attn_params


@torch.no_grad()
def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		int     ,							__lowercase:		Tuple     ,							__lowercase:		List[str]     ,							__lowercase:		List[str]     ,							__lowercase:		int=None):



  '''simple docstring'''






  __lowerCAmelCase			=				Repository(__lowercase     ,							clone_from=F'''google/{pytorch_dump_folder_path}''')
  repo.git_pull()

  if config_path is not None:
    __lowerCAmelCase			=				OwlViTConfig.from_pretrained(__lowercase)
  else:
    __lowerCAmelCase			=				OwlViTConfig()

  __lowerCAmelCase			=				OwlViTModel(__lowercase).eval()
  __lowerCAmelCase			=				OwlViTForObjectDetection(__lowercase).eval()

  copy_text_model_and_projection(__lowercase     ,							__lowercase)
  copy_vision_model_and_projection(__lowercase     ,							__lowercase)
  __lowerCAmelCase			=				pt_backbone.logit_scale
  copy_flax_attn_params(__lowercase     ,							__lowercase)

  __lowerCAmelCase			=				hf_backbone
  copy_class_merge_token(__lowercase     ,							__lowercase)
  copy_class_box_heads(__lowercase     ,							__lowercase)

  # 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=__lowercase     ,							tokenizer=__lowercase)
  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__":
      __lowercase													=			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."""
      )
      __lowercase													=			parser.parse_args()

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

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

      # Convert CLIP backbone
      __lowercase						,       __lowercase						,       __lowercase													=			convert_clip_backbone(flax_params, torch_config)

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







def 					__SCREAMING_SNAKE_CASE  (							__lowercase:		list):



  '''simple docstring'''






  if len(__lowercase) <= 1:
    return data
  else:
    return [
        *quick_sort([e for e in data[1:] if e <= data[0]]),
        data[0],
        *quick_sort([e for e in data[1:] if e > data[0]]),
    ]


if __name__ == "__main__":
      import doctest

      doctest.testmod()
