






"""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		:  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 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Optional[Any]		, __snake_case      :  Any=""		, __snake_case      :  str="/"   ):





  '''simple docstring'''





  lowercase							   =	[]

  for k, v in params.items():
    lowercase							   =	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 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  List[str]   ):





  '''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 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  int		, __snake_case      :  Dict   ):





  '''simple docstring'''





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

  lowercase							   =	pt_attn_layer.out_proj.weight
  lowercase							   =	pt_attn_layer.out_proj.bias

  lowercase							   =	q_proj
  lowercase							   =	q_proj_bias

  lowercase							   =	k_proj
  lowercase							   =	k_proj_bias

  lowercase							   =	v_proj
  lowercase							   =	v_proj_bias

  lowercase							   =	out_proj_weights
  lowercase							   =	out_proj_bias
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Optional[Any]		, __snake_case      :  Optional[Any]   ):





  '''simple docstring'''





  copy_linear(hf_mlp.fca		, pt_mlp.c_fc   )
  copy_linear(hf_mlp.fca		, pt_mlp.c_proj   )
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Optional[Any]		, __snake_case      :  Any   ):





  '''simple docstring'''





  lowercase							   =	pt_linear.weight
  lowercase							   =	pt_linear.bias
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Tuple		, __snake_case      :  Tuple   ):





  '''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			(   __snake_case      :  int		, __snake_case      :  Optional[int]   ):





  '''simple docstring'''





  for hf_layer, pt_layer in zip(__snake_case		, __snake_case   ):
    copy_layer(__snake_case		, __snake_case   )
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  int		, __snake_case      :  Any   ):





  '''simple docstring'''





  lowercase							   =	pt_model.token_embedding.weight
  lowercase							   =	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			(   __snake_case      :  Union[str, Any]		, __snake_case      :  List[Any]   ):





  '''simple docstring'''





  lowercase							   =	pt_model.text_projection.data.T

  # copy text encoder
  copy_encoder(hf_model.text_model		, __snake_case   )
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Optional[Any]		, __snake_case      :  Optional[int]   ):





  '''simple docstring'''





  lowercase							   =	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
  lowercase							   =	pt_model.visual.conva.weight.data
  lowercase							   =	pt_model.visual.class_embedding
  lowercase							   =	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			(   __snake_case      :  Union[str, Any]		, __snake_case      :  str   ):





  '''simple docstring'''





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

  lowercase							   =	torch.from_numpy(flax_class_token_params['scale']   )
  lowercase							   =	torch.from_numpy(flax_class_token_params['bias']   )
  lowercase							   =	nn.Parameter(__snake_case   )
  lowercase							   =	nn.Parameter(__snake_case   )
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  List[str]		, __snake_case      :  Any   ):





  '''simple docstring'''





  lowercase							   =	hf_model.state_dict()
  lowercase							   =	{}

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

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

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

    lowercase							   =	nn.Parameter(torch.from_numpy(__snake_case   )   )

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

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

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

    lowercase							   =	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 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  int		, __snake_case      :  List[str]   ):





  '''simple docstring'''





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

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

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

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

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

    # Copy flax CLIP attn params to HF PyTorch params
    lowercase							   =	torch.from_numpy(__snake_case   )
    hf_backbone.state_dict()[torch_key].copy_(__snake_case   )
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  str   ):





  '''simple docstring'''





  lowercase							   =	{}
  lowercase							   =	[]

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

      processed_attn_layers.append(__snake_case   )
      lowercase							   =	params[base + 'out.weight'].shape[-1]
      lowercase							   =	params[base + 'out.weight'].reshape(__snake_case		, __snake_case   ).T
      lowercase							   =	params[base + 'out.bias']
    else:
      lowercase							   =	v
  return new_params
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  Union[str, Any]		, __snake_case      :  Any   ):





  '''simple docstring'''





  lowercase							   =	CLIP(**__snake_case   )
  torch_model.eval()
  lowercase							   =	torch_model.state_dict()

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

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

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

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

    lowercase							   =	v

  lowercase							   =	_convert_attn_layers(__snake_case   )
  new_torch_params.update(__snake_case   )
  lowercase							   =	{}

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

  return torch_clip_params, torch_model, attn_params
@torch.no_grad()
def 						_SCREAMING_SNAKE_CASE			(   __snake_case      :  str		, __snake_case      :  List[str]		, __snake_case      :  int		, __snake_case      :  Optional[Any]		, __snake_case      :  Optional[int]=None   ):





  '''simple docstring'''





  lowercase							   =	Repository(__snake_case		, clone_from=f'google/{pytorch_dump_folder_path}'   )
  repo.git_pull()

  if config_path is not None:
    lowercase							   =	OwlViTConfig.from_pretrained(__snake_case   )
  else:
    lowercase							   =	OwlViTConfig()

  lowercase							   =	OwlViTModel(__snake_case   ).eval()
  lowercase							   =	OwlViTForObjectDetection(__snake_case   ).eval()

  copy_text_model_and_projection(__snake_case		, __snake_case   )
  copy_vision_model_and_projection(__snake_case		, __snake_case   )
  lowercase							   =	pt_backbone.logit_scale
  copy_flax_attn_params(__snake_case		, __snake_case   )

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

  # Initialize processor
  lowercase							   =	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		:  List[Any]						=		parser.parse_args()

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

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










"""simple docstring"""



import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor

class 	a					(       unittest.TestCase     ):







    def __init__(      self							,		_lowerCamelCase							,		_lowerCamelCase=7							,		_lowerCamelCase=3							,		_lowerCamelCase=1_8							,		_lowerCamelCase=3_0							,		_lowerCamelCase=4_0_0							,		_lowerCamelCase=True							,		_lowerCamelCase=None							,		_lowerCamelCase=True							,		_lowerCamelCase=[0.5, 0.5, 0.5]							,		_lowerCamelCase=[0.5, 0.5, 0.5]							,		):
      lowercase							   =	size if size is not None else {'height': 1_8, 'width': 1_8}
      lowercase							   =	parent
      lowercase							   =	batch_size
      lowercase							   =	num_channels
      lowercase							   =	image_size
      lowercase							   =	min_resolution
      lowercase							   =	max_resolution
      lowercase							   =	do_resize
      lowercase							   =	size
      lowercase							   =	do_normalize
      lowercase							   =	image_mean
      lowercase							   =	image_std







    def       UpperCamelCase_				(      self	):
      return {
          "image_mean": self.image_mean,
          "image_std": self.image_std,
          "do_normalize": self.do_normalize,
          "do_resize": self.do_resize,
          "size": self.size,
      }


@require_torch
@require_vision
class 	a					(       a_,    unittest.TestCase     ):
    UpperCAmelCase_			:					int     					=ViTImageProcessor if is_vision_available() else None







    def       UpperCamelCase_				(      self	):
      lowercase							   =	ViTImageProcessingTester(self	)







    @property
    def       UpperCamelCase_				(      self	):
      return self.image_processor_tester.prepare_image_processor_dict()







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.image_processing_class(**self.image_processor_dict	)
      self.assertTrue(hasattr(_lowerCamelCase							,		'image_mean'	)	)
      self.assertTrue(hasattr(_lowerCamelCase							,		'image_std'	)	)
      self.assertTrue(hasattr(_lowerCamelCase							,		'do_normalize'	)	)
      self.assertTrue(hasattr(_lowerCamelCase							,		'do_resize'	)	)
      self.assertTrue(hasattr(_lowerCamelCase							,		'size'	)	)







    def       UpperCamelCase_				(      self	):
      lowercase							   =	self.image_processing_class.from_dict(self.image_processor_dict	)
      self.assertEqual(image_processor.size							,		{'height': 1_8, 'width': 1_8}	)

      lowercase							   =	self.image_processing_class.from_dict(self.image_processor_dict							,		size=4_2	)
      self.assertEqual(image_processor.size							,		{'height': 4_2, 'width': 4_2}	)







    def       UpperCamelCase_				(      self	):
      pass







    def       UpperCamelCase_				(      self	):
      # Initialize image_processing
      lowercase							   =	self.image_processing_class(**self.image_processor_dict	)
      # create random PIL images
      lowercase							   =	prepare_image_inputs(self.image_processor_tester							,		equal_resolution=_lowerCamelCase	)
      for image in image_inputs:
        self.assertIsInstance(_lowerCamelCase							,		Image.Image	)

      # Test not batched input
      lowercase							   =	image_processing(image_inputs[0]							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)

      # Test batched
      lowercase							   =	image_processing(_lowerCamelCase							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)







    def       UpperCamelCase_				(      self	):
      # Initialize image_processing
      lowercase							   =	self.image_processing_class(**self.image_processor_dict	)
      # create random numpy tensors
      lowercase							   =	prepare_image_inputs(self.image_processor_tester							,		equal_resolution=_lowerCamelCase							,		numpify=_lowerCamelCase	)
      for image in image_inputs:
        self.assertIsInstance(_lowerCamelCase							,		np.ndarray	)

      # Test not batched input
      lowercase							   =	image_processing(image_inputs[0]							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)

      # Test batched
      lowercase							   =	image_processing(_lowerCamelCase							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)







    def       UpperCamelCase_				(      self	):
      # Initialize image_processing
      lowercase							   =	self.image_processing_class(**self.image_processor_dict	)
      # create random PyTorch tensors
      lowercase							   =	prepare_image_inputs(self.image_processor_tester							,		equal_resolution=_lowerCamelCase							,		torchify=_lowerCamelCase	)
      for image in image_inputs:
        self.assertIsInstance(_lowerCamelCase							,		torch.Tensor	)

      # Test not batched input
      lowercase							   =	image_processing(image_inputs[0]							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)

      # Test batched
      lowercase							   =	image_processing(_lowerCamelCase							,		return_tensors='pt'	).pixel_values
      self.assertEqual(
          encoded_images.shape							,		(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.size['height'],
              self.image_processor_tester.size['width'],
          )							,		)


