





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()
_lowerCAmelCase							:      Tuple						=	logging.get_logger(__name__)

def      __snake_case						(					_lowerCAmelCase : Optional[Any]	) ->						Any:
   A_					:				int          =			tf.image.resize(_lowerCAmelCase				,				(346, 346)	)
   A_					:				Dict          =			tf.image.crop_to_bounding_box(_lowerCAmelCase				,				(346 - 289) // 2				,				(346 - 289) // 2				,				289				,				289	)
   return image

def      __snake_case						(					) ->						Tuple:
   A_					:				str          =			EfficientNetConfig.from_pretrained("google/efficientnet-b7"	)
   A_					:				Optional[int]          =			289
   A_					:				Union[str, Any]          =			640
   A_					:				Dict          =			{"0": "LABEL_0", "1": "LABEL_1"}
   A_					:				List[Any]          =			{"LABEL_0": 0, "LABEL_1": 1}
   A_					:				List[str]          =			[]

   A_					:				List[Any]          =			BertConfig()
   A_					:				Any          =			AlignConfig.from_text_vision_configs(
       text_config=_lowerCAmelCase				,				vision_config=_lowerCAmelCase				,				projection_dim=640	)
   return config

def      __snake_case						(					) ->						List[str]:
   A_					:				str          =			"http://images.cocodataset.org/val2017/000000039769.jpg"
   A_					:				Optional[Any]          =			Image.open(requests.get(_lowerCAmelCase				,				stream=_lowerCAmelCase	).raw	)
   return im

def      __snake_case						(					) ->						Optional[int]:
   A_					:				List[str]          =			EfficientNetImageProcessor(
       do_center_crop=_lowerCAmelCase				,				rescale_factor=1 / 127.5				,				rescale_offset=_lowerCAmelCase				,				do_normalize=_lowerCAmelCase				,				include_top=_lowerCAmelCase				,				resample=Image.BILINEAR				,				)
   A_					:				Optional[Any]          =			BertTokenizer.from_pretrained("bert-base-uncased"	)
   A_					:				Tuple          =			64
   A_					:				List[Any]          =			AlignProcessor(image_processor=_lowerCAmelCase				,				tokenizer=_lowerCAmelCase	)
   return processor

def      __snake_case						(					_lowerCAmelCase : str	) ->						Optional[Any]:
   # EfficientNet image encoder
   A_					:				Tuple          =			[v.split("_"	)[0].split("block"	)[1] for v in original_param_names if v.startswith("block"	)]
   A_					:				Union[str, Any]          =			list(set(_lowerCAmelCase	)	)
   A_					:				Union[str, Any]          =			sorted(_lowerCAmelCase	)
   A_					:				Any          =			len(_lowerCAmelCase	)
   A_					:				str          =			{b: str(_lowerCAmelCase	) for b, i in zip(_lowerCAmelCase				,				range(_lowerCAmelCase	)	)}

   A_					:				Dict          =			[]
   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:
      A_					:				Optional[int]          =			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")	)

   A_					:				List[str]          =			{}
   for item in rename_keys:
      if item[0] in original_param_names:
         A_					:				Optional[int]          =			"vision_model." + item[1]

    # BERT text encoder
   A_					:				Tuple          =			[]
   A_					:				List[Any]          =			"tf_bert_model/bert"
   A_					:				int          =			"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:
         A_					:				Tuple          =			item[1]
   return key_mapping

def      __snake_case						(					_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : Tuple				,				_lowerCAmelCase : Tuple	) ->						Union[str, Any]:
   list(hf_params.keys()	)

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

      A_					:				Optional[Any]          =			key_mapping[key]
      if "_conv" in key and "kernel" in key:
         A_					:				Optional[Any]          =			torch.from_numpy(_lowerCAmelCase	).permute(3				,				2				,				0				,				1	)
      elif "embeddings" in key:
         A_					:				Optional[int]          =			torch.from_numpy(_lowerCAmelCase	)
      elif "depthwise_kernel" in key:
         A_					:				Tuple          =			torch.from_numpy(_lowerCAmelCase	).permute(2				,				3				,				0				,				1	)
      elif "kernel" in key:
         A_					:				str          =			torch.from_numpy(np.transpose(_lowerCAmelCase	)	)
      elif "temperature" in key:
         A_					:				Any          =			value
      elif "bn/gamma" or "bn/beta" in key:
         A_					:				List[Any]          =			torch.from_numpy(np.transpose(_lowerCAmelCase	)	).squeeze()
      else:
         A_					:				Dict          =			torch.from_numpy(_lowerCAmelCase	)

      # Replace HF parameters with original TF model parameters
      hf_params[hf_key].copy_(_lowerCAmelCase	)

@torch.no_grad()
def      __snake_case						(					_lowerCAmelCase : List[Any]				,				_lowerCAmelCase : Any				,				_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : int	) ->						List[Any]:
   A_					:				Dict          =			64
   A_					:				Dict          =			Tokenizer(_lowerCAmelCase	)
   A_					:				Union[str, Any]          =			align.Align("efficientnet-b7"				,				"bert-base"				,				640				,				_lowerCAmelCase				,				tok.get_vocab_size()	)
   original_model.compile()
   original_model.load_weights(_lowerCAmelCase	)

   A_					:				Any          =			original_model.trainable_variables
   A_					:				Any          =			original_model.non_trainable_variables
   A_					:				List[str]          =			{param.name: param.numpy() for param in tf_params}
   for param in tf_non_train_params:
      A_					:				List[Any]          =			param.numpy()
   A_					:				int          =			list(tf_params.keys()	)

   # Load HuggingFace model
   A_					:				Optional[int]          =			get_align_config()
   A_					:				List[Any]          =			AlignModel(_lowerCAmelCase	).eval()
   A_					:				Dict          =			hf_model.state_dict()

   # Create src-to-dst parameter name mapping dictionary
   print("Converting parameters..."	)
   A_					:				Dict          =			rename_keys(_lowerCAmelCase	)
   replace_params(_lowerCAmelCase				,				_lowerCAmelCase				,				_lowerCAmelCase	)

   # Initialize processor
   A_					:				int          =			get_processor()
   A_					:				Any          =			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():
      A_					:				Optional[Any]          =			hf_model(**_lowerCAmelCase	)

   A_					:				List[str]          =			outputs.image_embeds.detach().numpy()
   A_					:				Tuple          =			outputs.text_embeds.detach().numpy()

   # Original model inference
   A_					:				List[str]          =			False
   A_					:				List[str]          =			EfficientNetImageProcessor(
       do_center_crop=_lowerCAmelCase				,				do_rescale=_lowerCAmelCase				,				do_normalize=_lowerCAmelCase				,				include_top=_lowerCAmelCase				,				resample=Image.BILINEAR				,				)
   A_					:				List[str]          =			tf_image_processor(images=prepare_img()				,				return_tensors="tf"				,				data_format="channels_last"	)["pixel_values"]
   A_					:				Dict          =			tok(tf.constant(["A picture of a cat"]	)	)

   A_					:				List[Any]          =			original_model.image_encoder(_lowerCAmelCase				,				training=_lowerCAmelCase	)
   A_					:				List[Any]          =			original_model.text_encoder(_lowerCAmelCase				,				training=_lowerCAmelCase	)

   A_					:				Any          =			tf.nn.la_normalize(_lowerCAmelCase				,				axis=-1	)
   A_					:				Tuple          =			tf.nn.la_normalize(_lowerCAmelCase				,				axis=-1	)

   # Check whether original and HF model outputs match  -> np.allclose
   if not np.allclose(_lowerCAmelCase				,				_lowerCAmelCase				,				atol=1e-3	):
      raise ValueError("The predicted image features are not the same."	)
   if not np.allclose(_lowerCAmelCase				,				_lowerCAmelCase				,				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(_lowerCAmelCase	):
         os.mkdir(_lowerCAmelCase	)
      # Save converted model and image processor
      hf_model.save_pretrained(_lowerCAmelCase	)
      processor.save_pretrained(_lowerCAmelCase	)

   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__":
 _lowerCAmelCase							:      List[Any]						=	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''')

 _lowerCAmelCase							:      List[str]						=	parser.parse_args()
 convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)






def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : int = 1000	) ->						str:
   if (not isinstance(_lowerCAmelCase				,				_lowerCAmelCase	)) or (digit_position <= 0):
      raise ValueError("Digit position must be a positive integer"	)
   elif (not isinstance(_lowerCAmelCase				,				_lowerCAmelCase	)) or (precision < 0):
      raise ValueError("Precision must be a nonnegative integer"	)

   # compute an approximation of (16 ** (n - 1)) * pi whose fractional part is mostly
   # accurate
   A_					:				List[str]          =			(
       4 * _subsum(_lowerCAmelCase				,				1				,				_lowerCAmelCase	)
       - 2 * _subsum(_lowerCAmelCase				,				4				,				_lowerCAmelCase	)
       - _subsum(_lowerCAmelCase				,				5				,				_lowerCAmelCase	)
       - _subsum(_lowerCAmelCase				,				6				,				_lowerCAmelCase	)
   )

   # return the first hex digit of the fractional part of the result
   return hex(int((sum_result % 1) * 16	)	)[2:]

def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : int				,				_lowerCAmelCase : int	) ->						float:
   A_					:				int          =			0.0
   for sum_index in range(digit_pos_to_extract + precision	):
      A_					:				Optional[int]          =			8 * sum_index + denominator_addend
      if sum_index < digit_pos_to_extract:
         # if the exponential term is an integer and we mod it by the denominator
         # before dividing, only the integer part of the sum will change;
         # the fractional part will not
         A_					:				Optional[Any]          =			pow(
             16				,				digit_pos_to_extract - 1 - sum_index				,				_lowerCAmelCase	)
      else:
         A_					:				Tuple          =			pow(16				,				digit_pos_to_extract - 1 - sum_index	)
      total += exponential_term / denominator
   return total


if __name__ == "__main__":
 import doctest

 doctest.testmod()
