


import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase			:					Dict         =							logging.get_logger(__name__)
def 				UpperCamelCase_	(					__snake_case      :    int  )->   int:




						'''simple docstring'''
						__magic_name__					=      tf.image.resize(__snake_case						,      (3_46, 3_46)  )
						__magic_name__					=      tf.image.crop_to_bounding_box(__snake_case						,      (3_46 - 2_89) // 2						,      (3_46 - 2_89) // 2						,      2_89						,      2_89  )
						return image
def 				UpperCamelCase_	(					)->   int:




						'''simple docstring'''
						__magic_name__					=      EfficientNetConfig.from_pretrained('google/efficientnet-b7'  )
						__magic_name__					=      2_89
						__magic_name__					=      6_40
						__magic_name__					=      {'0': 'LABEL_0', '1': 'LABEL_1'}
						__magic_name__					=      {'LABEL_0': 0, 'LABEL_1': 1}
						__magic_name__					=      []

						__magic_name__					=      BertConfig()
						__magic_name__					=      AlignConfig.from_text_vision_configs(
						    text_config=__snake_case						,      vision_config=__snake_case						,      projection_dim=6_40  )
						return config
def 				UpperCamelCase_	(					)->   Optional[Any]:




						'''simple docstring'''
						__magic_name__					=      'http://images.cocodataset.org/val2017/000000039769.jpg'
						__magic_name__					=      Image.open(requests.get(__snake_case						,      stream=__snake_case  ).raw  )
						return im
def 				UpperCamelCase_	(					)->   Optional[Any]:




						'''simple docstring'''
						__magic_name__					=      EfficientNetImageProcessor(
						    do_center_crop=__snake_case						,      rescale_factor=1 / 127.5						,      rescale_offset=__snake_case						,      do_normalize=__snake_case						,      include_top=__snake_case						,      resample=Image.BILINEAR						,      )
						__magic_name__					=      BertTokenizer.from_pretrained('bert-base-uncased'  )
						__magic_name__					=      64
						__magic_name__					=      AlignProcessor(image_processor=__snake_case						,      tokenizer=__snake_case  )
						return processor
def 				UpperCamelCase_	(					__snake_case      :    Any  )->   Tuple:




						'''simple docstring'''
						# EfficientNet image encoder
						__magic_name__					=      [v.split('_'  )[0].split('block'  )[1] for v in original_param_names if v.startswith('block'  )]
						__magic_name__					=      list(set(__snake_case  )  )
						__magic_name__					=      sorted(__snake_case  )
						__magic_name__					=      len(__snake_case  )
						__magic_name__					=      {b: str(__snake_case  ) for b, i in zip(__snake_case						,      range(__snake_case  )  )}

						__magic_name__					=      []
						rename_keys.append(('stem_conv/kernel:0', 'embeddings.convolution.weight')  )
						rename_keys.append(('stem_bn/gamma:0', 'embeddings.batchnorm.weight')  )
						rename_keys.append(('stem_bn/beta:0', 'embeddings.batchnorm.bias')  )
						rename_keys.append(('stem_bn/moving_mean:0', 'embeddings.batchnorm.running_mean')  )
						rename_keys.append(('stem_bn/moving_variance:0', 'embeddings.batchnorm.running_var')  )

						for b in block_names:
												__magic_name__					=      block_name_mapping[b]
												rename_keys.append((F'''block{b}_expand_conv/kernel:0''', F'''encoder.blocks.{hf_b}.expansion.expand_conv.weight''')  )
												rename_keys.append((F'''block{b}_expand_bn/gamma:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.weight''')  )
												rename_keys.append((F'''block{b}_expand_bn/beta:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.bias''')  )
												rename_keys.append(
												    (F'''block{b}_expand_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.running_mean''')  )
												rename_keys.append(
												    (F'''block{b}_expand_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.expansion.expand_bn.running_var''')  )
												rename_keys.append(
												    (F'''block{b}_dwconv/depthwise_kernel:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight''')  )
												rename_keys.append((F'''block{b}_bn/gamma:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight''')  )
												rename_keys.append((F'''block{b}_bn/beta:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias''')  )
												rename_keys.append(
												    (F'''block{b}_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean''')  )
												rename_keys.append(
												    (F'''block{b}_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var''')  )

												rename_keys.append((F'''block{b}_se_reduce/kernel:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.reduce.weight''')  )
												rename_keys.append((F'''block{b}_se_reduce/bias:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.reduce.bias''')  )
												rename_keys.append((F'''block{b}_se_expand/kernel:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.expand.weight''')  )
												rename_keys.append((F'''block{b}_se_expand/bias:0''', F'''encoder.blocks.{hf_b}.squeeze_excite.expand.bias''')  )
												rename_keys.append(
												    (F'''block{b}_project_conv/kernel:0''', F'''encoder.blocks.{hf_b}.projection.project_conv.weight''')  )
												rename_keys.append((F'''block{b}_project_bn/gamma:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.weight''')  )
												rename_keys.append((F'''block{b}_project_bn/beta:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.bias''')  )
												rename_keys.append(
												    (F'''block{b}_project_bn/moving_mean:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.running_mean''')  )
												rename_keys.append(
												    (F'''block{b}_project_bn/moving_variance:0''', F'''encoder.blocks.{hf_b}.projection.project_bn.running_var''')  )

						__magic_name__					=      {}
						for item in rename_keys:
												if item[0] in original_param_names:
																		__magic_name__					=      'vision_model.' + item[1]

    # BERT text encoder
						__magic_name__					=      []
						__magic_name__					=      'tf_bert_model/bert'
						__magic_name__					=      'text_model'
						for i in range(12  ):
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/query/kernel:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.query.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/query/bias:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.query.bias''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/key/kernel:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.key.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/key/bias:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.key.bias''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/value/kernel:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.value.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/self/value/bias:0''',
												        F'''{new}.encoder.layer.{i}.attention.self.value.bias''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/output/dense/kernel:0''',
												        F'''{new}.encoder.layer.{i}.attention.output.dense.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/output/dense/bias:0''',
												        F'''{new}.encoder.layer.{i}.attention.output.dense.bias''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0''',
												        F'''{new}.encoder.layer.{i}.attention.output.LayerNorm.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0''',
												        F'''{new}.encoder.layer.{i}.attention.output.LayerNorm.bias''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/intermediate/dense/kernel:0''',
												        F'''{new}.encoder.layer.{i}.intermediate.dense.weight''',
												    )  )
												rename_keys.append(
												    (
												        F'''{old}/encoder/layer_._{i}/intermediate/dense/bias:0''',
												        F'''{new}.encoder.layer.{i}.intermediate.dense.bias''',
												    )  )
												rename_keys.append(
												    (F'''{old}/encoder/layer_._{i}/output/dense/kernel:0''', F'''{new}.encoder.layer.{i}.output.dense.weight''')  )
												rename_keys.append(
												    (F'''{old}/encoder/layer_._{i}/output/dense/bias:0''', F'''{new}.encoder.layer.{i}.output.dense.bias''')  )
												rename_keys.append(
												    (F'''{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0''', F'''{new}.encoder.layer.{i}.output.LayerNorm.weight''')  )
												rename_keys.append(
												    (F'''{old}/encoder/layer_._{i}/output/LayerNorm/beta:0''', F'''{new}.encoder.layer.{i}.output.LayerNorm.bias''')  )

						rename_keys.append((F'''{old}/embeddings/word_embeddings/weight:0''', F'''{new}.embeddings.word_embeddings.weight''')  )
						rename_keys.append(
						    (F'''{old}/embeddings/position_embeddings/embeddings:0''', F'''{new}.embeddings.position_embeddings.weight''')  )
						rename_keys.append(
						    (F'''{old}/embeddings/token_type_embeddings/embeddings:0''', F'''{new}.embeddings.token_type_embeddings.weight''')  )
						rename_keys.append((F'''{old}/embeddings/LayerNorm/gamma:0''', F'''{new}.embeddings.LayerNorm.weight''')  )
						rename_keys.append((F'''{old}/embeddings/LayerNorm/beta:0''', F'''{new}.embeddings.LayerNorm.bias''')  )

						rename_keys.append((F'''{old}/pooler/dense/kernel:0''', F'''{new}.pooler.dense.weight''')  )
						rename_keys.append((F'''{old}/pooler/dense/bias:0''', F'''{new}.pooler.dense.bias''')  )
						rename_keys.append(('dense/kernel:0', 'text_projection.weight')  )
						rename_keys.append(('dense/bias:0', 'text_projection.bias')  )
						rename_keys.append(('dense/bias:0', 'text_projection.bias')  )
						rename_keys.append(('temperature:0', 'temperature')  )

						for item in rename_keys:
												if item[0] in original_param_names:
																		__magic_name__					=      item[1]
						return key_mapping
def 				UpperCamelCase_	(					__snake_case      :    Optional[Any]						,      __snake_case      :    Optional[int]						,      __snake_case      :    Optional[Any]  )->   str:




						'''simple docstring'''
						list(hf_params.keys()  )

						for key, value in tf_params.items():
												if key not in key_mapping:
																		continue

												__magic_name__					=      key_mapping[key]
												if "_conv" in key and "kernel" in key:
																		__magic_name__					=      torch.from_numpy(__snake_case  ).permute(3						,      2						,      0						,      1  )
												elif "embeddings" in key:
																		__magic_name__					=      torch.from_numpy(__snake_case  )
												elif "depthwise_kernel" in key:
																		__magic_name__					=      torch.from_numpy(__snake_case  ).permute(2						,      3						,      0						,      1  )
												elif "kernel" in key:
																		__magic_name__					=      torch.from_numpy(np.transpose(__snake_case  )  )
												elif "temperature" in key:
																		__magic_name__					=      value
												elif "bn/gamma" or "bn/beta" in key:
																		__magic_name__					=      torch.from_numpy(np.transpose(__snake_case  )  ).squeeze()
												else:
																		__magic_name__					=      torch.from_numpy(__snake_case  )

												# Replace HF parameters with original TF model parameters
												hf_params[hf_key].copy_(__snake_case  )
@torch.no_grad()
def 				UpperCamelCase_	(					__snake_case      :    Optional[Any]						,      __snake_case      :    Optional[int]						,      __snake_case      :    List[str]						,      __snake_case      :    Union[str, Any]  )->   int:




						'''simple docstring'''
						__magic_name__					=      64
						__magic_name__					=      Tokenizer(__snake_case  )
						__magic_name__					=      align.Align('efficientnet-b7'						,      'bert-base'						,      6_40						,      __snake_case						,      tok.get_vocab_size()  )
						original_model.compile()
						original_model.load_weights(__snake_case  )

						__magic_name__					=      original_model.trainable_variables
						__magic_name__					=      original_model.non_trainable_variables
						__magic_name__					=      {param.name: param.numpy() for param in tf_params}
						for param in tf_non_train_params:
												__magic_name__					=      param.numpy()
						__magic_name__					=      list(tf_params.keys()  )

						# Load HuggingFace model
						__magic_name__					=      get_align_config()
						__magic_name__					=      AlignModel(__snake_case  ).eval()
						__magic_name__					=      hf_model.state_dict()

						# Create src-to-dst parameter name mapping dictionary
						print('Converting parameters...'  )
						__magic_name__					=      rename_keys(__snake_case  )
						replace_params(__snake_case						,      __snake_case						,      __snake_case  )

						# Initialize processor
						__magic_name__					=      get_processor()
						__magic_name__					=      processor(
						    images=prepare_img()						,      text='A picture of a cat'						,      padding='max_length'						,      max_length=64						,      return_tensors='pt'  )

						# HF model inference
						hf_model.eval()
						with torch.no_grad():
												__magic_name__					=      hf_model(**__snake_case  )

						__magic_name__					=      outputs.image_embeds.detach().numpy()
						__magic_name__					=      outputs.text_embeds.detach().numpy()

						# Original model inference
						__magic_name__					=      False
						__magic_name__					=      EfficientNetImageProcessor(
						    do_center_crop=__snake_case						,      do_rescale=__snake_case						,      do_normalize=__snake_case						,      include_top=__snake_case						,      resample=Image.BILINEAR						,      )
						__magic_name__					=      tf_image_processor(images=prepare_img()						,      return_tensors='tf'						,      data_format='channels_last'  )['pixel_values']
						__magic_name__					=      tok(tf.constant(['A picture of a cat']  )  )

						__magic_name__					=      original_model.image_encoder(__snake_case						,      training=__snake_case  )
						__magic_name__					=      original_model.text_encoder(__snake_case						,      training=__snake_case  )

						__magic_name__					=      tf.nn.la_normalize(__snake_case						,      axis=-1  )
						__magic_name__					=      tf.nn.la_normalize(__snake_case						,      axis=-1  )

						# Check whether original and HF model outputs match  -> np.allclose
						if not np.allclose(__snake_case						,      __snake_case						,      atol=1E-3  ):
												raise ValueError('The predicted image features are not the same.'  )
						if not np.allclose(__snake_case						,      __snake_case						,      atol=1E-3  ):
												raise ValueError('The predicted text features are not the same.'  )
						print('Model outputs match!'  )

						if save_model:
												# Create folder to save model
												if not os.path.isdir(__snake_case  ):
																		os.mkdir(__snake_case  )
												# Save converted model and image processor
												hf_model.save_pretrained(__snake_case  )
												processor.save_pretrained(__snake_case  )

						if push_to_hub:
												# Push model and image processor to hub
												print('Pushing converted ALIGN to the hub...'  )
												processor.push_to_hub('align-base'  )
												hf_model.push_to_hub('align-base'  )


if __name__ == "__main__":
					UpperCamelCase			:					List[str]         =							argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--checkpoint_path''',
					    default='''./weights/model-weights''',
					    type=str,
					    help='''Path to the pretrained TF ALIGN checkpoint.''',
					)
					parser.add_argument(
					    '''--pytorch_dump_folder_path''',
					    default='''hf_model''',
					    type=str,
					    help='''Path to the output PyTorch model directory.''',
					)
					parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
					parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')

					UpperCamelCase			:					int         =							parser.parse_args()
					convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)





from collections import OrderedDict
from dataclasses import fields
from typing import Any, Tuple

import numpy as np

from .import_utils import is_torch_available
def 				UpperCamelCase_	(					__snake_case      :    Dict  )->   str:




						'''simple docstring'''
						if is_torch_available():
												import torch

												if isinstance(__snake_case						,      torch.Tensor  ):
																		return True

						return isinstance(__snake_case						,      np.ndarray  )




class      __snake_case				(   _UpperCAmelCase):
					def 	a			(						self       :						Optional[Any]  )							->					Optional[Any]:



											'''simple docstring'''





											__magic_name__					=      fields(self  )

											# Safety and consistency checks
											if not len(_lowerCAmelCase  ):
																	raise ValueError(F'''{self.__class__.__name__} has no fields.'''  )

											__magic_name__					=      getattr(self	,      class_fields[0].name  )
											__magic_name__					=      all(getattr(self	,      field.name  ) is None for field in class_fields[1:]  )

											if other_fields_are_none and isinstance(_lowerCAmelCase	,      _lowerCAmelCase  ):
																	for key, value in first_field.items():
																							__magic_name__					=      value
											else:
																	for field in class_fields:
																							__magic_name__					=      getattr(self	,      field.name  )
																							if v is not None:
																													__magic_name__					=      v
					def __delitem__(						self       :						Dict	,      *_lowerCAmelCase       :						List[str]	,      **_lowerCAmelCase       :						List[str]  )							->					List[str]:



											'''simple docstring'''





											raise Exception(F'''You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.'''  )
					def 	a			(						self       :						Optional[int]	,      *_lowerCAmelCase       :						Optional[Any]	,      **_lowerCAmelCase       :						Optional[Any]  )							->					str:



											'''simple docstring'''





											raise Exception(F'''You cannot use ``setdefault`` on a {self.__class__.__name__} instance.'''  )
					def 	a			(						self       :						Union[str, Any]	,      *_lowerCAmelCase       :						List[str]	,      **_lowerCAmelCase       :						Any  )							->					str:



											'''simple docstring'''





											raise Exception(F'''You cannot use ``pop`` on a {self.__class__.__name__} instance.'''  )
					def 	a			(						self       :						Any	,      *_lowerCAmelCase       :						str	,      **_lowerCAmelCase       :						Optional[int]  )							->					List[Any]:



											'''simple docstring'''





											raise Exception(F'''You cannot use ``update`` on a {self.__class__.__name__} instance.'''  )
					def __getitem__(						self       :						Optional[Any]	,      _lowerCAmelCase       :						Any  )							->					Tuple:



											'''simple docstring'''





											if isinstance(_lowerCAmelCase	,      _lowerCAmelCase  ):
																	__magic_name__					=      dict(self.items()  )
																	return inner_dict[k]
											else:
																	return self.to_tuple()[k]
					def __setattr__(						self       :						str	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						Optional[int]  )							->					int:



											'''simple docstring'''





											if name in self.keys() and value is not None:
																	# Don't call self.__setitem__ to avoid recursion errors
																	super().__setitem__(_lowerCAmelCase	,      _lowerCAmelCase  )
											super().__setattr__(_lowerCAmelCase	,      _lowerCAmelCase  )
					def __setitem__(						self       :						Optional[Any]	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						List[str]  )							->					Optional[int]:



											'''simple docstring'''





											super().__setitem__(_lowerCAmelCase	,      _lowerCAmelCase  )
											# Don't call self.__setattr__ to avoid recursion errors
											super().__setattr__(_lowerCAmelCase	,      _lowerCAmelCase  )







					def 	a			(						self       :						str  )							->					Tuple[Any]:



											'''simple docstring'''





											return tuple(self[k] for k in self.keys()  )

