


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




def 			__lowercase	(						_UpperCamelCase     )    ->Tuple:







      """simple docstring"""

      lowercase			:   Optional[int]					   =		tf.image.resize(_UpperCamelCase,    (346, 346)     )
      lowercase			:   Any					   =		tf.image.crop_to_bounding_box(_UpperCamelCase,    (346 - 289) // 2,    (346 - 289) // 2,    289,    289     )
      return image




def 			__lowercase	(						)    ->List[str]:







      """simple docstring"""

      lowercase			:   Tuple					   =		EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''     )
      lowercase			:   Optional[int]					   =		289
      lowercase			:   Optional[Any]					   =		640
      lowercase			:   Any					   =		{'''0''': '''LABEL_0''', '''1''': '''LABEL_1'''}
      lowercase			:   Optional[Any]					   =		{'''LABEL_0''': 0, '''LABEL_1''': 1}
      lowercase			:   Union[str, Any]					   =		[]

      lowercase			:   Union[str, Any]					   =		BertConfig()
      lowercase			:   Any					   =		AlignConfig.from_text_vision_configs(
          text_config=_UpperCamelCase,    vision_config=_UpperCamelCase,    projection_dim=640     )
      return config




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







      """simple docstring"""

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




def 			__lowercase	(						)    ->Optional[int]:







      """simple docstring"""

      lowercase			:   Any					   =		EfficientNetImageProcessor(
          do_center_crop=_UpperCamelCase,    rescale_factor=1 / 1_2_7.5,    rescale_offset=_UpperCamelCase,    do_normalize=_UpperCamelCase,    include_top=_UpperCamelCase,    resample=Image.BILINEAR,    )
      lowercase			:   Any					   =		BertTokenizer.from_pretrained('''bert-base-uncased'''     )
      lowercase			:   int					   =		64
      lowercase			:   List[Any]					   =		AlignProcessor(image_processor=_UpperCamelCase,    tokenizer=_UpperCamelCase     )
      return processor




def 			__lowercase	(						_UpperCamelCase     )    ->int:







      """simple docstring"""

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

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

      lowercase			:   Union[str, Any]					   =		{}
      for item in rename_keys:
            if item[0] in original_param_names:
                  lowercase			:   Dict					   =		'''vision_model.''' + item[1]

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




def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )    ->Any:







      """simple docstring"""

      list(hf_params.keys()     )

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

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

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




@torch.no_grad()
def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )    ->Dict:







      """simple docstring"""

      lowercase			:   str					   =		64
      lowercase			:   str					   =		Tokenizer(_UpperCamelCase     )
      lowercase			:   Optional[int]					   =		align.Align('''efficientnet-b7''',    '''bert-base''',    640,    _UpperCamelCase,    tok.get_vocab_size()     )
      original_model.compile()
      original_model.load_weights(_UpperCamelCase     )

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

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

      # Create src-to-dst parameter name mapping dictionary
      print('''Converting parameters...'''     )
      lowercase			:   List[str]					   =		rename_keys(_UpperCamelCase     )
      replace_params(_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )

      # Initialize processor
      lowercase			:   Union[str, Any]					   =		get_processor()
      lowercase			:   Optional[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():
            lowercase			:   Dict					   =		hf_model(**_UpperCamelCase     )

      lowercase			:   int					   =		outputs.image_embeds.detach().numpy()
      lowercase			:   Any					   =		outputs.text_embeds.detach().numpy()

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

      lowercase			:   Optional[int]					   =		original_model.image_encoder(_UpperCamelCase,    training=_UpperCamelCase     )
      lowercase			:   Optional[int]					   =		original_model.text_encoder(_UpperCamelCase,    training=_UpperCamelCase     )

      lowercase			:   Union[str, Any]					   =		tf.nn.la_normalize(_UpperCamelCase,    axis=-1     )
      lowercase			:   List[str]					   =		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   				=      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   				=      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 			__lowercase	(						_UpperCamelCase     )    ->list[int]:







      """simple docstring"""

      assert (
          isinstance(_UpperCamelCase,    _UpperCamelCase     ) and mask > 0
      ), f"""mask needs to be positive integer, your input {mask}"""
      lowercase			:   Optional[int]					   =		[]
      lowercase			:   Optional[int]					   =		mask

      while submask:
            all_submasks.append(_UpperCamelCase     )
            lowercase			:   Dict					   =		(submask - 1) & mask

      return all_submasks


if __name__ == "__main__":
  import doctest

  doctest.testmod()

