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()
__snake_case			      =    logging.get_logger(__name__)





def   lowerCAmelCase_ (				__lowerCAmelCase						)->  Tuple:






      '''simple docstring'''





      UpperCAmelCase    : Optional[Any]      =tf.image.resize(__lowerCAmelCase		,					(3_46, 3_46)						)
      UpperCAmelCase    : List[str]      =tf.image.crop_to_bounding_box(__lowerCAmelCase		,					(3_46 - 2_89) // 2		,					(3_46 - 2_89) // 2		,					2_89		,					2_89						)
      return image





def   lowerCAmelCase_ (				)->  Tuple:






      '''simple docstring'''





      UpperCAmelCase    : Union[str, Any]      =EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''						)
      UpperCAmelCase    : Optional[int]      =2_89
      UpperCAmelCase    : Tuple      =6_40
      UpperCAmelCase    : List[str]      ={'''0''': '''LABEL_0''', '''1''': '''LABEL_1'''}
      UpperCAmelCase    : List[str]      ={'''LABEL_0''': 0, '''LABEL_1''': 1}
      UpperCAmelCase    : int      =[]

      UpperCAmelCase    : Dict      =BertConfig()
      UpperCAmelCase    : Optional[int]      =AlignConfig.from_text_vision_configs(
          text_config=__lowerCAmelCase		,					vision_config=__lowerCAmelCase		,					projection_dim=6_40						)
      return config





def   lowerCAmelCase_ (				)->  List[str]:






      '''simple docstring'''





      UpperCAmelCase    : int      ='''http://images.cocodataset.org/val2017/000000039769.jpg'''
      UpperCAmelCase    : Any      =Image.open(requests.get(__lowerCAmelCase		,					stream=__lowerCAmelCase						).raw						)
      return im





def   lowerCAmelCase_ (				)->  Optional[Any]:






      '''simple docstring'''





      UpperCAmelCase    : Optional[Any]      =EfficientNetImageProcessor(
          do_center_crop=__lowerCAmelCase		,					rescale_factor=1 / 127.5		,					rescale_offset=__lowerCAmelCase		,					do_normalize=__lowerCAmelCase		,					include_top=__lowerCAmelCase		,					resample=Image.BILINEAR		,					)
      UpperCAmelCase    : List[Any]      =BertTokenizer.from_pretrained('''bert-base-uncased'''						)
      UpperCAmelCase    : Optional[Any]      =64
      UpperCAmelCase    : Dict      =AlignProcessor(image_processor=__lowerCAmelCase		,					tokenizer=__lowerCAmelCase						)
      return processor





def   lowerCAmelCase_ (				__lowerCAmelCase						)->  Union[str, Any]:






      '''simple docstring'''





      UpperCAmelCase    : Union[str, Any]      =[v.split('''_'''						)[0].split('''block'''						)[1] for v in original_param_names if v.startswith('''block'''						)]
      UpperCAmelCase    : Dict      =list(set(__lowerCAmelCase						)						)
      UpperCAmelCase    : Tuple      =sorted(__lowerCAmelCase						)
      UpperCAmelCase    : List[Any]      =len(__lowerCAmelCase						)
      UpperCAmelCase    : str      ={b: str(__lowerCAmelCase						) for b, i in zip(__lowerCAmelCase		,					range(__lowerCAmelCase						)						)}

      UpperCAmelCase    : Any      =[]
      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:
            UpperCAmelCase    : Any      =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''')						)

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

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





def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase						)->  List[str]:






      '''simple docstring'''





      list(hf_params.keys()						)

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

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

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





@torch.no_grad()
def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase						)->  Any:






      '''simple docstring'''





      UpperCAmelCase    : Tuple      =64
      UpperCAmelCase    : Any      =Tokenizer(__lowerCAmelCase						)
      UpperCAmelCase    : List[Any]      =align.Align('''efficientnet-b7'''		,					'''bert-base'''		,					6_40		,					__lowerCAmelCase		,					tok.get_vocab_size()						)
      original_model.compile()
      original_model.load_weights(__lowerCAmelCase						)

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

      # Load HuggingFace model
      UpperCAmelCase    : List[str]      =get_align_config()
      UpperCAmelCase    : Union[str, Any]      =AlignModel(__lowerCAmelCase						).eval()
      UpperCAmelCase    : List[Any]      =hf_model.state_dict()

      # Create src-to-dst parameter name mapping dictionary
      print('''Converting parameters...'''						)
      UpperCAmelCase    : str      =rename_keys(__lowerCAmelCase						)
      replace_params(__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase						)

      # Initialize processor
      UpperCAmelCase    : List[Any]      =get_processor()
      UpperCAmelCase    : Dict      =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():
            UpperCAmelCase    : Tuple      =hf_model(**__lowerCAmelCase						)

      UpperCAmelCase    : int      =outputs.image_embeds.detach().numpy()
      UpperCAmelCase    : Union[str, Any]      =outputs.text_embeds.detach().numpy()

      # Original model inference
      UpperCAmelCase    : Union[str, Any]      =False
      UpperCAmelCase    : Optional[int]      =EfficientNetImageProcessor(
          do_center_crop=__lowerCAmelCase		,					do_rescale=__lowerCAmelCase		,					do_normalize=__lowerCAmelCase		,					include_top=__lowerCAmelCase		,					resample=Image.BILINEAR		,					)
      UpperCAmelCase    : Tuple      =tf_image_processor(images=prepare_img()		,					return_tensors='''tf'''		,					data_format='''channels_last'''						)['''pixel_values''']
      UpperCAmelCase    : Union[str, Any]      =tok(tf.constant(['''A picture of a cat''']						)						)

      UpperCAmelCase    : List[str]      =original_model.image_encoder(__lowerCAmelCase		,					training=__lowerCAmelCase						)
      UpperCAmelCase    : Dict      =original_model.text_encoder(__lowerCAmelCase		,					training=__lowerCAmelCase						)

      UpperCAmelCase    : List[str]      =tf.nn.la_normalize(__lowerCAmelCase		,					axis=-1						)
      UpperCAmelCase    : Union[str, Any]      =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__":
     __snake_case			      =    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''')

     __snake_case			      =    parser.parse_args()
     convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


__snake_case			      =    {
    '''configuration_esm''': ['''ESM_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''EsmConfig'''],
    '''tokenization_esm''': ['''EsmTokenizer'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __snake_case			      =    [
         '''ESM_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''EsmForMaskedLM''',
         '''EsmForSequenceClassification''',
         '''EsmForTokenClassification''',
         '''EsmModel''',
         '''EsmPreTrainedModel''',
     ]
     __snake_case			      =    ['''EsmForProteinFolding''', '''EsmFoldPreTrainedModel''']

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __snake_case			      =    [
         '''TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFEsmForMaskedLM''',
         '''TFEsmForSequenceClassification''',
         '''TFEsmForTokenClassification''',
         '''TFEsmModel''',
         '''TFEsmPreTrainedModel''',
     ]

if TYPE_CHECKING:
     from .configuration_esm import ESM_PRETRAINED_CONFIG_ARCHIVE_MAP, EsmConfig
     from .tokenization_esm import EsmTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_esm import (
              ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
              EsmForMaskedLM,
              EsmForSequenceClassification,
              EsmForTokenClassification,
              EsmModel,
              EsmPreTrainedModel,
          )
          from .modeling_esmfold import EsmFoldPreTrainedModel, EsmForProteinFolding

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_esm import (
              TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
              TFEsmModel,
              TFEsmPreTrainedModel,
          )


else:
     import sys

     __snake_case			      =    _LazyModule(__name__, globals()['''__file__'''], _import_structure)
