
'''simple docstring'''





import collections
import inspect
import unittest

from transformers import SwinConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import SwinBackbone, SwinForImageClassification, SwinForMaskedImageModeling, SwinModel
 from transformers.models.swin.modeling_swin import SWIN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor


class 			UpperCamelCase__	:






 def __init__(      self	:List[str]						,					_A	:str						,					_A	:Optional[Any]=13						,					_A	:Tuple=32						,					_A	:Dict=2						,					_A	:str=3						,					_A	:Union[str, Any]=16						,					_A	:Optional[Any]=[1, 2, 1]						,					_A	:Tuple=[2, 2, 4]						,					_A	:Dict=2						,					_A	:Optional[int]=2.0						,					_A	:int=True						,					_A	:List[Any]=0.0						,					_A	:Dict=0.0						,					_A	:int=0.1						,					_A	:Optional[Any]="gelu"						,					_A	:Union[str, Any]=False						,					_A	:Optional[Any]=True						,					_A	:Union[str, Any]=0.02						,					_A	:Optional[Any]=1E-5						,					_A	:Union[str, Any]=True						,					_A	:Optional[int]=None						,					_A	:int=True						,					_A	:Optional[Any]=10						,					_A	:List[Any]=8						,					_A	:Optional[Any]=["stage1", "stage2"]						,					_A	:Any=[1, 2]						,					)					->      List[Any]:







       '''simple docstring'''



       __A													=      parent
       __A													=      batch_size
       __A													=      image_size
       __A													=      patch_size
       __A													=      num_channels
       __A													=      embed_dim
       __A													=      depths
       __A													=      num_heads
       __A													=      window_size
       __A													=      mlp_ratio
       __A													=      qkv_bias
       __A													=      hidden_dropout_prob
       __A													=      attention_probs_dropout_prob
       __A													=      drop_path_rate
       __A													=      hidden_act
       __A													=      use_absolute_embeddings
       __A													=      patch_norm
       __A													=      layer_norm_eps
       __A													=      initializer_range
       __A													=      is_training
       __A													=      scope
       __A													=      use_labels
       __A													=      type_sequence_label_size
       __A													=      encoder_stride
       __A													=      out_features
       __A													=      out_indices






 def 						lowercase_    (      self	:Tuple      )					->      List[Any]:







       '''simple docstring'''



       __A													=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

       __A													=      None
       if self.use_labels:
             __A													=      ids_tensor([self.batch_size]						,					self.type_sequence_label_size      )

       __A													=      self.get_config()

       return config, pixel_values, labels






 def 						lowercase_    (      self	:str      )					->      List[str]:







       '''simple docstring'''



       return SwinConfig(
           image_size=self.image_size						,					patch_size=self.patch_size						,					num_channels=self.num_channels						,					embed_dim=self.embed_dim						,					depths=self.depths						,					num_heads=self.num_heads						,					window_size=self.window_size						,					mlp_ratio=self.mlp_ratio						,					qkv_bias=self.qkv_bias						,					hidden_dropout_prob=self.hidden_dropout_prob						,					attention_probs_dropout_prob=self.attention_probs_dropout_prob						,					drop_path_rate=self.drop_path_rate						,					hidden_act=self.hidden_act						,					use_absolute_embeddings=self.use_absolute_embeddings						,					path_norm=self.patch_norm						,					layer_norm_eps=self.layer_norm_eps						,					initializer_range=self.initializer_range						,					encoder_stride=self.encoder_stride						,					out_features=self.out_features						,					out_indices=self.out_indices						,					)






 def 						lowercase_    (      self	:Tuple						,					_A	:Optional[int]						,					_A	:Dict						,					_A	:List[Any]      )					->      str:







       '''simple docstring'''



       __A													=      SwinModel(config=_A      )
       model.to(_A      )
       model.eval()
       __A													=      model(_A      )

       __A													=      ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths      ) - 1))
       __A													=      int(config.embed_dim * 2 ** (len(config.depths      ) - 1)      )

       self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, expected_seq_len, expected_dim)      )






 def 						lowercase_    (      self	:Optional[int]						,					_A	:Optional[Any]						,					_A	:List[str]						,					_A	:List[Any]      )					->      Dict:







       '''simple docstring'''



       __A													=      SwinBackbone(config=_A      )
       model.to(_A      )
       model.eval()
       __A													=      model(_A      )

       # verify hidden states
       self.parent.assertEqual(len(result.feature_maps      )						,					len(config.out_features      )      )
       self.parent.assertListEqual(list(result.feature_maps[0].shape      )						,					[self.batch_size, model.channels[0], 16, 16]      )

       # verify channels
       self.parent.assertEqual(len(model.channels      )						,					len(config.out_features      )      )

       # verify backbone works with out_features=None
       __A													=      None
       __A													=      SwinBackbone(config=_A      )
       model.to(_A      )
       model.eval()
       __A													=      model(_A      )

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps      )						,					1      )
       self.parent.assertListEqual(list(result.feature_maps[0].shape      )						,					[self.batch_size, model.channels[-1], 4, 4]      )

       # verify channels
       self.parent.assertEqual(len(model.channels      )						,					1      )






 def 						lowercase_    (      self	:int						,					_A	:Optional[int]						,					_A	:int						,					_A	:Union[str, Any]      )					->      Optional[int]:







       '''simple docstring'''



       __A													=      SwinForMaskedImageModeling(config=_A      )
       model.to(_A      )
       model.eval()
       __A													=      model(_A      )
       self.parent.assertEqual(
           result.logits.shape						,					(self.batch_size, self.num_channels, self.image_size, self.image_size)      )

       # test greyscale images
       __A													=      1
       __A													=      SwinForMaskedImageModeling(_A      )
       model.to(_A      )
       model.eval()

       __A													=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
       __A													=      model(_A      )
       self.parent.assertEqual(result.logits.shape						,					(self.batch_size, 1, self.image_size, self.image_size)      )






 def 						lowercase_    (      self	:List[Any]						,					_A	:Union[str, Any]						,					_A	:List[Any]						,					_A	:Tuple      )					->      Any:







       '''simple docstring'''



       __A													=      self.type_sequence_label_size
       __A													=      SwinForImageClassification(_A      )
       model.to(_A      )
       model.eval()
       __A													=      model(_A						,					labels=_A      )
       self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)      )

       # test greyscale images
       __A													=      1
       __A													=      SwinForImageClassification(_A      )
       model.to(_A      )
       model.eval()

       __A													=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
       __A													=      model(_A      )
       self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.type_sequence_label_size)      )






 def 						lowercase_    (      self	:Dict      )					->      Dict:







       '''simple docstring'''



       __A													=      self.prepare_config_and_inputs()
       (
           (
           __A
       )						,   (
           __A
       )						,   (
           __A
       )						,   
       )													=      config_and_inputs
       __A													=      {'pixel_values': pixel_values}
       return config, inputs_dict


@require_torch
class 			UpperCamelCase__	(     SCREAMING_SNAKE_CASE   ,  SCREAMING_SNAKE_CASE   ,  unittest.TestCase):
 UpperCAmelCase__     :     str    					=       (
     (
         SwinModel,
         SwinBackbone,
         SwinForImageClassification,
         SwinForMaskedImageModeling,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase__     :     Optional[Any]    					=       (
     {'feature-extraction': SwinModel, 'image-classification': SwinForImageClassification}
     if is_torch_available()
     else {}
 )
 UpperCAmelCase__     :     Dict    					=       True

 UpperCAmelCase__     :     List[Any]    					=       False
 UpperCAmelCase__     :     Tuple    					=       False
 UpperCAmelCase__     :     str    					=       False






 def 						lowercase_    (      self	:Optional[Any]      )					->      List[Any]:







       '''simple docstring'''



       __A													=      SwinModelTester(self      )
       __A													=      ConfigTester(self						,					config_class=_A						,					embed_dim=37      )






 def 						lowercase_    (      self	:List[Any]      )					->      str:







       '''simple docstring'''



       self.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()






 def 						lowercase_    (      self	:Union[str, Any]      )					->      Optional[Any]:







       '''simple docstring'''



       return






 def 						lowercase_    (      self	:Union[str, Any]      )					->      List[Any]:







       '''simple docstring'''



       __A													=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_A      )






 @unittest.skip(reason='Got `CUDA error: misaligned address` with PyTorch 2.0.0.'      )
 def 						lowercase_    (      self	:int      )					->      Any:







       '''simple docstring'''



       pass






 def 						lowercase_    (      self	:Dict      )					->      Dict:







       '''simple docstring'''



       super().test_training_gradient_checkpointing()






 def 						lowercase_    (      self	:List[str]      )					->      Optional[int]:







       '''simple docstring'''



       __A													=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*_A      )






 def 						lowercase_    (      self	:Optional[int]      )					->      Optional[Any]:







       '''simple docstring'''



       __A													=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*_A      )






 def 						lowercase_    (      self	:Optional[Any]      )					->      Dict:







       '''simple docstring'''



       __A													=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_A      )






 @unittest.skip(reason='Swin does not use inputs_embeds'      )
 def 						lowercase_    (      self	:Union[str, Any]      )					->      Dict:







       '''simple docstring'''



       pass






 @unittest.skip(reason='Swin Transformer does not use feedforward chunking'      )
 def 						lowercase_    (      self	:str      )					->      Any:







       '''simple docstring'''



       pass






 def 						lowercase_    (      self	:List[str]      )					->      Dict:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             __A													=      model_class(_A      )
             self.assertIsInstance(model.get_input_embeddings()						,					(nn.Module)      )
             __A													=      model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_A						,					nn.Linear      )      )






 def 						lowercase_    (      self	:List[str]      )					->      List[str]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             __A													=      model_class(_A      )
             __A													=      inspect.signature(model.forward      )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             __A													=      [*signature.parameters.keys()]

             __A													=      ['pixel_values']
             self.assertListEqual(arg_names[:1]						,					_A      )






 def 						lowercase_    (      self	:Dict      )					->      Union[str, Any]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()
       __A													=      True

       for model_class in self.all_model_classes:
             __A													=      True
             __A													=      False
             __A													=      True
             __A													=      model_class(_A      )
             model.to(_A      )
             model.eval()
             with torch.no_grad():
                   __A													=      model(**self._prepare_for_class(_A						,					_A      )      )
             __A													=      outputs.attentions
             __A													=      len(self.model_tester.depths      )
             self.assertEqual(len(_A      )						,					_A      )

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             __A													=      True
             __A													=      config.window_size**2
             __A													=      model_class(_A      )
             model.to(_A      )
             model.eval()
             with torch.no_grad():
                   __A													=      model(**self._prepare_for_class(_A						,					_A      )      )
             __A													=      outputs.attentions
             self.assertEqual(len(_A      )						,					_A      )

             self.assertListEqual(
                 list(attentions[0].shape[-3:]      )						,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared]						,					)
             __A													=      len(_A      )

             # Check attention is always last and order is fine
             __A													=      True
             __A													=      True
             __A													=      model_class(_A      )
             model.to(_A      )
             model.eval()
             with torch.no_grad():
                   __A													=      model(**self._prepare_for_class(_A						,					_A      )      )

             # also another +1 for reshaped_hidden_states
             __A													=      1 if model_class.__name__ == 'SwinBackbone' else 2
             self.assertEqual(out_len + added_hidden_states						,					len(_A      )      )

             __A													=      outputs.attentions

             self.assertEqual(len(_A      )						,					_A      )

             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]      )						,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared]						,					)






 def 						lowercase_    (      self	:int						,					_A	:int						,					_A	:Dict						,					_A	:Optional[Any]						,					_A	:Dict      )					->      Dict:







       '''simple docstring'''



       __A													=      model_class(_A      )
       model.to(_A      )
       model.eval()

       with torch.no_grad():
             __A													=      model(**self._prepare_for_class(_A						,					_A      )      )

       __A													=      outputs.hidden_states

       __A													=      getattr(
           self.model_tester						,					'expected_num_hidden_layers'						,					len(self.model_tester.depths      ) + 1      )
       self.assertEqual(len(_A      )						,					_A      )

       # Swin has a different seq_length
       __A													=      (
           config.patch_size
           if isinstance(config.patch_size						,					collections.abc.Iterable      )
           else (config.patch_size, config.patch_size)
       )

       __A													=      (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

       self.assertListEqual(
           list(hidden_states[0].shape[-2:]      )						,					[num_patches, self.model_tester.embed_dim]						,					)

       if not model_class.__name__ == "SwinBackbone":
             __A													=      outputs.reshaped_hidden_states
             self.assertEqual(len(_A      )						,					_A      )

             __A						,   __A						,   __A						,   __A													=      reshaped_hidden_states[0].shape
             __A													=      (
                 reshaped_hidden_states[0].view(_A						,					_A						,					height * width      ).permute(0						,					2						,					1      )
             )
             self.assertListEqual(
                 list(reshaped_hidden_states.shape[-2:]      )						,					[num_patches, self.model_tester.embed_dim]						,					)






 def 						lowercase_    (      self	:int      )					->      Optional[Any]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()

       __A													=      (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size						,					collections.abc.Iterable      )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )

       for model_class in self.all_model_classes:
             __A													=      True
             self.check_hidden_states_output(_A						,					_A						,					_A						,					_A      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             __A													=      True

             self.check_hidden_states_output(_A						,					_A						,					_A						,					_A      )






 def 						lowercase_    (      self	:int      )					->      Union[str, Any]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()
       __A													=      3

       __A													=      (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size						,					collections.abc.Iterable      )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )
       __A													=      (
           config.patch_size
           if isinstance(config.patch_size						,					collections.abc.Iterable      )
           else (config.patch_size, config.patch_size)
       )

       __A													=      image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
       __A													=      image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

       for model_class in self.all_model_classes:
             __A													=      True
             self.check_hidden_states_output(_A						,					_A						,					_A						,					(padded_height, padded_width)      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             __A													=      True
             self.check_hidden_states_output(_A						,					_A						,					_A						,					(padded_height, padded_width)      )






 @slow
 def 						lowercase_    (      self	:Optional[Any]      )					->      Dict:







       '''simple docstring'''



       for model_name in SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A													=      SwinModel.from_pretrained(_A      )
             self.assertIsNotNone(_A      )






 def 						lowercase_    (      self	:str      )					->      List[str]:







       '''simple docstring'''



       __A						,   __A													=      self.model_tester.prepare_config_and_inputs_for_common()

       __A													=      _config_zero_init(_A      )
       for model_class in self.all_model_classes:
             __A													=      model_class(config=_A      )
             for name, param in model.named_parameters():
                   if "embeddings" not in name and param.requires_grad:
                         self.assertIn(
                             ((param.data.mean() * 1E9).round() / 1E9).item()						,					[0.0, 1.0]						,					msg=F'Parameter {name} of model {model_class} seems not properly initialized'						,					)


@require_vision
@require_torch
class 			UpperCamelCase__	(     unittest.TestCase):






 @cached_property
 def 						lowercase_    (      self	:Any      )					->      Any:







       '''simple docstring'''



       return (
           AutoImageProcessor.from_pretrained('microsoft/swin-tiny-patch4-window7-224'      )
           if is_vision_available()
           else None
       )






 @slow
 def 						lowercase_    (      self	:List[Any]      )					->      Any:







       '''simple docstring'''



       __A													=      SwinForImageClassification.from_pretrained('microsoft/swin-tiny-patch4-window7-224'      ).to(_A      )
       __A													=      self.default_image_processor

       __A													=      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
       __A													=      image_processor(images=_A						,					return_tensors='pt'      ).to(_A      )

       # forward pass
       with torch.no_grad():
             __A													=      model(**_A      )

       # verify the logits
       __A													=      torch.Size((1, 1_000)      )
       self.assertEqual(outputs.logits.shape						,					_A      )
       __A													=      torch.tensor([-0.0_948, -0.6_454, -0.0_921]      ).to(_A      )
       self.assertTrue(torch.allclose(outputs.logits[0, :3]						,					_A						,					atol=1E-4      )      )


@require_torch
class 			UpperCamelCase__	(     unittest.TestCase   ,  SCREAMING_SNAKE_CASE):
 UpperCAmelCase__     :     int    					=       (SwinBackbone,) if is_torch_available() else ()
 UpperCAmelCase__     :     Optional[int]    					=       SwinConfig






 def 						lowercase_    (      self	:str      )					->      Union[str, Any]:







       '''simple docstring'''



       __A													=      SwinModelTester(self      )









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


_lowerCamelCase							:  List[Any]     =				{
    "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      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     Dict=""							,					UpperCAmelCase				:     Optional[Any]="/"     )    ->		Optional[int]:
  UpperCAmelCase			:				str										=					[]

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

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

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

def      a__	(						UpperCAmelCase				:     Union[str, Any]							,					UpperCAmelCase				:     List[Any]     )    ->		str:
  UpperCAmelCase					,		UpperCAmelCase					,		UpperCAmelCase			:				List[Any]										=					pt_attn_layer.in_proj_weight.chunk(3							,					dim=0     )
  UpperCAmelCase					,		UpperCAmelCase					,		UpperCAmelCase			:				Tuple										=					pt_attn_layer.in_proj_bias.chunk(3							,					dim=0     )

  UpperCAmelCase			:				Optional[Any]										=					pt_attn_layer.out_proj.weight
  UpperCAmelCase			:				Tuple										=					pt_attn_layer.out_proj.bias

  UpperCAmelCase			:				List[str]										=					q_proj
  UpperCAmelCase			:				Optional[int]										=					q_proj_bias

  UpperCAmelCase			:				Any										=					k_proj
  UpperCAmelCase			:				Optional[Any]										=					k_proj_bias

  UpperCAmelCase			:				Optional[Any]										=					v_proj
  UpperCAmelCase			:				Optional[Any]										=					v_proj_bias

  UpperCAmelCase			:				str										=					out_proj_weights
  UpperCAmelCase			:				List[str]										=					out_proj_bias

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Dict     )    ->		Optional[int]:
  copy_linear(hf_mlp.fca							,					pt_mlp.c_fc     )
  copy_linear(hf_mlp.fca							,					pt_mlp.c_proj     )

def      a__	(						UpperCAmelCase				:     Optional[int]							,					UpperCAmelCase				:     str     )    ->		Dict:
  UpperCAmelCase			:				Dict										=					pt_linear.weight
  UpperCAmelCase			:				List[Any]										=					pt_linear.bias

def      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     int     )    ->		Tuple:
  # copy layer norms
  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      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     Union[str, Any]     )    ->		int:
  for hf_layer, pt_layer in zip(UpperCAmelCase							,					UpperCAmelCase     ):
    copy_layer(UpperCAmelCase							,					UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     Dict     )    ->		Union[str, Any]:
  # copy  embeds
  UpperCAmelCase			:				List[Any]										=					pt_model.token_embedding.weight
  UpperCAmelCase			:				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      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     Tuple     )    ->		str:
  # copy projection
  UpperCAmelCase			:				int										=					pt_model.text_projection.data.T

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

def      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     List[str]     )    ->		List[Any]:
  # copy projection
  UpperCAmelCase			:				List[Any]										=					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			:				Dict										=					pt_model.visual.conva.weight.data
  UpperCAmelCase			:				Any										=					pt_model.visual.class_embedding
  UpperCAmelCase			:				str										=					pt_model.visual.positional_embedding.data

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

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Any     )    ->		Union[str, Any]:
  UpperCAmelCase			:				str										=					flatten_nested_dict(flax_params['''backbone''']['''merged_class_token''']     )

  UpperCAmelCase			:				str										=					torch.from_numpy(flax_class_token_params['''scale''']     )
  UpperCAmelCase			:				Union[str, Any]										=					torch.from_numpy(flax_class_token_params['''bias''']     )
  UpperCAmelCase			:				Union[str, Any]										=					nn.Parameter(UpperCAmelCase     )
  UpperCAmelCase			:				int										=					nn.Parameter(UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     Any							,					UpperCAmelCase				:     Optional[Any]     )    ->		List[Any]:
  UpperCAmelCase			:				int										=					hf_model.state_dict()
  UpperCAmelCase			:				Any										=					{}

  # Rename class prediction head flax params to pytorch HF
  UpperCAmelCase			:				Optional[Any]										=					flatten_nested_dict(flax_params['''class_head''']     )

  for flax_key, v in flax_class_params.items():
    UpperCAmelCase			:				List[str]										=					flax_key.replace('''/'''							,					'''.'''     )
    UpperCAmelCase			:				Dict										=					torch_key.replace('''.kernel'''							,					'''.weight'''     )
    UpperCAmelCase			:				Optional[int]										=					torch_key.replace('''Dense_0'''							,					'''dense0'''     )
    UpperCAmelCase			:				List[Any]										=					'''class_head.''' + torch_key

    if "weight" in torch_key and v.ndim == 2:
      UpperCAmelCase			:				Union[str, Any]										=					v.T

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

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

  for flax_key, v in flax_box_params.items():
    UpperCAmelCase			:				Any										=					flax_key.replace('''/'''							,					'''.'''     )
    UpperCAmelCase			:				Any										=					torch_key.replace('''.kernel'''							,					'''.weight'''     )
    UpperCAmelCase			:				List[Any]										=					torch_key.replace('''_'''							,					''''''     ).lower()
    UpperCAmelCase			:				int										=					'''box_head.''' + torch_key

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

    UpperCAmelCase			:				int										=					nn.Parameter(torch.from_numpy(UpperCAmelCase     )     )

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

def      a__	(						UpperCAmelCase				:     int							,					UpperCAmelCase				:     List[Any]     )    ->		List[Any]:
  for k, v in flax_attn_params.items():
    if k.startswith('''transformer'''     ):
      UpperCAmelCase			:				Optional[int]										=					k.replace('''transformer.resblocks'''							,					'''text_model.encoder.layers'''     )
    else:
      UpperCAmelCase			:				Tuple										=					k.replace('''visual.transformer.resblocks'''							,					'''vision_model.encoder.layers'''     )

    UpperCAmelCase			:				List[Any]										=					torch_key.replace('''attn'''							,					'''self_attn'''     )
    UpperCAmelCase			:				Union[str, Any]										=					torch_key.replace('''key'''							,					'''k_proj'''     )
    UpperCAmelCase			:				int										=					torch_key.replace('''value'''							,					'''v_proj'''     )
    UpperCAmelCase			:				Dict										=					torch_key.replace('''query'''							,					'''q_proj'''     )
    UpperCAmelCase			:				int										=					torch_key.replace('''out'''							,					'''out_proj'''     )

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

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

    # Copy flax CLIP attn params to HF PyTorch params
    UpperCAmelCase			:				List[str]										=					torch.from_numpy(UpperCAmelCase     )
    hf_backbone.state_dict()[torch_key].copy_(UpperCAmelCase     )

def      a__	(						UpperCAmelCase				:     Union[str, Any]     )    ->		List[str]:
  UpperCAmelCase			:				Optional[Any]										=					{}
  UpperCAmelCase			:				Dict										=					[]

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

      processed_attn_layers.append(UpperCAmelCase     )
      UpperCAmelCase			:				List[str]										=					params[base + '''out.weight'''].shape[-1]
      UpperCAmelCase			:				Tuple										=					params[base + '''out.weight'''].reshape(UpperCAmelCase							,					UpperCAmelCase     ).T
      UpperCAmelCase			:				Tuple										=					params[base + '''out.bias''']
    else:
      UpperCAmelCase			:				Optional[int]										=					v
  return new_params

def      a__	(						UpperCAmelCase				:     Tuple							,					UpperCAmelCase				:     Optional[Any]     )    ->		int:
  UpperCAmelCase			:				Optional[int]										=					CLIP(**UpperCAmelCase     )
  torch_model.eval()
  UpperCAmelCase			:				Optional[Any]										=					torch_model.state_dict()

  UpperCAmelCase			:				List[str]										=					flatten_nested_dict(flax_params['''backbone''']['''clip''']     )
  UpperCAmelCase			:				Any										=					{}

  for flax_key, v in flax_clip_params.items():
    UpperCAmelCase			:				Optional[int]										=					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			:				List[str]										=					torch_key[5:]

    UpperCAmelCase			:				Dict										=					torch_key.replace('''text_projection.kernel'''							,					'''text_projection'''     )
    UpperCAmelCase			:				Tuple										=					torch_key.replace('''visual.proj.kernel'''							,					'''visual.proj'''     )
    UpperCAmelCase			:				Dict										=					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			:				Dict										=					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			:				Dict										=					v.T

    UpperCAmelCase			:				Union[str, Any]										=					v

  UpperCAmelCase			:				Dict										=					_convert_attn_layers(UpperCAmelCase     )
  new_torch_params.update(UpperCAmelCase     )
  UpperCAmelCase			:				int										=					{}

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

  return torch_clip_params, torch_model, attn_params

@torch.no_grad()
def      a__	(						UpperCAmelCase				:     str							,					UpperCAmelCase				:     Union[str, Any]							,					UpperCAmelCase				:     Optional[int]							,					UpperCAmelCase				:     List[str]							,					UpperCAmelCase				:     Optional[Any]=None     )    ->		Optional[int]:
  UpperCAmelCase			:				Optional[int]										=					Repository(UpperCAmelCase							,					clone_from=f'''google/{pytorch_dump_folder_path}'''     )
  repo.git_pull()

  if config_path is not None:
    UpperCAmelCase			:				str										=					OwlViTConfig.from_pretrained(UpperCAmelCase     )
  else:
    UpperCAmelCase			:				List[str]										=					OwlViTConfig()

  UpperCAmelCase			:				List[str]										=					OwlViTModel(UpperCAmelCase     ).eval()
  UpperCAmelCase			:				Optional[int]										=					OwlViTForObjectDetection(UpperCAmelCase     ).eval()

  copy_text_model_and_projection(UpperCAmelCase							,					UpperCAmelCase     )
  copy_vision_model_and_projection(UpperCAmelCase							,					UpperCAmelCase     )
  UpperCAmelCase			:				Tuple										=					pt_backbone.logit_scale
  copy_flax_attn_params(UpperCAmelCase							,					UpperCAmelCase     )

  UpperCAmelCase			:				int										=					hf_backbone
  copy_class_merge_token(UpperCAmelCase							,					UpperCAmelCase     )
  copy_class_box_heads(UpperCAmelCase							,					UpperCAmelCase     )

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

  # Initialize image processor
  UpperCAmelCase			:				List[str]										=					OwlViTImageProcessor(
      size=config.vision_config.image_size							,					crop_size=config.vision_config.image_size     )
  # Initialize tokenizer
  UpperCAmelCase			:				Union[str, Any]										=					CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''							,					pad_token='''!'''							,					model_max_length=16     )

  # Initialize processor
  UpperCAmelCase			:				str										=					OwlViTProcessor(image_processor=UpperCAmelCase							,					tokenizer=UpperCAmelCase     )
  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__":
     _lowerCamelCase							:  List[str]     =				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."
     )
     _lowerCamelCase							:  List[Any]     =				parser.parse_args()

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

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

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





