





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()
a__			:    List[Any]			      =		logging.get_logger(__name__)

def     UpperCAmelCase_						(						_UpperCAmelCase       :Tuple							)  ->							Union[str, Any]:
     '''simple docstring'''


     A_					      =					tf.image.resize(_UpperCAmelCase     ,	(346, 346)							)
     A_					      =					tf.image.crop_to_bounding_box(_UpperCAmelCase     ,	(346 - 289) // 2     ,	(346 - 289) // 2     ,	289     ,	289							)
     return image

def     UpperCAmelCase_						(						)  ->							str:
     '''simple docstring'''


     A_					      =					EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''							)
     A_					      =					289
     A_					      =					640
     A_					      =					{"0": "LABEL_0", "1": "LABEL_1"}
     A_					      =					{"LABEL_0": 0, "LABEL_1": 1}
     A_					      =					[]

     A_					      =					BertConfig()
     A_					      =					AlignConfig.from_text_vision_configs(
         text_config=_UpperCAmelCase     ,	vision_config=_UpperCAmelCase     ,	projection_dim=640							)
     return config

def     UpperCAmelCase_						(						)  ->							str:
     '''simple docstring'''


     A_					      =					"http://images.cocodataset.org/val2017/000000039769.jpg"
     A_					      =					Image.open(requests.get(_UpperCAmelCase     ,	stream=_UpperCAmelCase							).raw							)
     return im

def     UpperCAmelCase_						(						)  ->							Optional[Any]:
     '''simple docstring'''


     A_					      =					EfficientNetImageProcessor(
         do_center_crop=_UpperCAmelCase     ,	rescale_factor=1 / 127.5     ,	rescale_offset=_UpperCAmelCase     ,	do_normalize=_UpperCAmelCase     ,	include_top=_UpperCAmelCase     ,	resample=Image.BILINEAR     ,	)
     A_					      =					BertTokenizer.from_pretrained('''bert-base-uncased'''							)
     A_					      =					64
     A_					      =					AlignProcessor(image_processor=_UpperCAmelCase     ,	tokenizer=_UpperCAmelCase							)
     return processor

def     UpperCAmelCase_						(						_UpperCAmelCase       :int							)  ->							Union[str, Any]:
     '''simple docstring'''


     A_					      =					[v.split('''_'''							)[0].split('''block'''							)[1] for v in original_param_names if v.startswith('''block'''							)]
     A_					      =					list(set(_UpperCAmelCase							)							)
     A_					      =					sorted(_UpperCAmelCase							)
     A_					      =					len(_UpperCAmelCase							)
     A_					      =					{b: str(_UpperCAmelCase							) for b, i in zip(_UpperCAmelCase     ,	range(_UpperCAmelCase							)							)}

     A_					      =					[]
     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_					      =					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_					      =					{}
     for item in rename_keys:
          if item[0] in original_param_names:
               A_					      =					"vision_model." + item[1]

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

def     UpperCAmelCase_						(						_UpperCAmelCase       :Optional[Any]     ,	_UpperCAmelCase       :int     ,	_UpperCAmelCase       :Tuple							)  ->							Tuple:
     '''simple docstring'''


     list(hf_params.keys()							)

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

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

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

@torch.no_grad()
def     UpperCAmelCase_						(						_UpperCAmelCase       :int     ,	_UpperCAmelCase       :Union[str, Any]     ,	_UpperCAmelCase       :str     ,	_UpperCAmelCase       :Any							)  ->							List[Any]:
     '''simple docstring'''


     A_					      =					64
     A_					      =					Tokenizer(_UpperCAmelCase							)
     A_					      =					align.Align('''efficientnet-b7'''     ,	'''bert-base'''     ,	640     ,	_UpperCAmelCase     ,	tok.get_vocab_size()							)
     original_model.compile()
     original_model.load_weights(_UpperCAmelCase							)

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

     # Load HuggingFace model
     A_					      =					get_align_config()
     A_					      =					AlignModel(_UpperCAmelCase							).eval()
     A_					      =					hf_model.state_dict()

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

     # Initialize processor
     A_					      =					get_processor()
     A_					      =					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_					      =					hf_model(**_UpperCAmelCase							)

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

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

     A_					      =					original_model.image_encoder(_UpperCAmelCase     ,	training=_UpperCAmelCase							)
     A_					      =					original_model.text_encoder(_UpperCAmelCase     ,	training=_UpperCAmelCase							)

     A_					      =					tf.nn.la_normalize(_UpperCAmelCase     ,	axis=-1							)
     A_					      =					tf.nn.la_normalize(_UpperCAmelCase     ,	axis=-1							)

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

     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__":
      a__			:    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')

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










from __future__ import annotations

def     UpperCAmelCase_						(						_UpperCAmelCase       :list							)  ->							int | float:
     '''simple docstring'''


     A_					      =					sorted(_UpperCAmelCase							)
     A_					      =					len(_UpperCAmelCase							)
     A_					      =					length >> 1
     return (
         (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
         if length % 2 == 0
         else sorted_list[mid_index]
     )

def     UpperCAmelCase_						(						)  ->							Optional[Any]:
     '''simple docstring'''


     import doctest

     doctest.testmod()


if __name__ == "__main__":
      main()



