




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

def       lowerCAmelCase_      (	lowerCamelCase:							Tuple    )       ->      List[str]:


      '''simple docstring'''




      lowerCamelCase__   	=					tf.image.resize(lowerCamelCase  ,					(3_46, 3_46)    )
      lowerCamelCase__   	=					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'''




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

      lowerCamelCase__   	=					BertConfig()
      lowerCamelCase__   	=					AlignConfig.from_text_vision_configs(
          text_config=lowerCamelCase  ,					vision_config=lowerCamelCase  ,					projection_dim=6_40    )
      return config

def       lowerCAmelCase_      (	)       ->      Optional[int]:


      '''simple docstring'''




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

def       lowerCAmelCase_      (	)       ->      List[str]:


      '''simple docstring'''




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

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


      '''simple docstring'''




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

      lowerCamelCase__   	=					[]
      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:
            lowerCamelCase__   	=					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""")    )

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

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

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


      '''simple docstring'''




      list(hf_params.keys()    )

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

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

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

@torch.no_grad()
def       lowerCAmelCase_      (	lowerCamelCase:							List[Any]  ,					lowerCamelCase:							Dict  ,					lowerCamelCase:							List[Any]  ,					lowerCamelCase:							int    )       ->      Union[str, Any]:


      '''simple docstring'''




      lowerCamelCase__   	=					64
      lowerCamelCase__   	=					Tokenizer(lowerCamelCase    )
      lowerCamelCase__   	=					align.Align('efficientnet-b7'  ,					'bert-base'  ,					6_40  ,					lowerCamelCase  ,					tok.get_vocab_size()    )
      original_model.compile()
      original_model.load_weights(lowerCamelCase    )

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

      # Load HuggingFace model
      lowerCamelCase__   	=					get_align_config()
      lowerCamelCase__   	=					AlignModel(lowerCamelCase    ).eval()
      lowerCamelCase__   	=					hf_model.state_dict()

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

      # Initialize processor
      lowerCamelCase__   	=					get_processor()
      lowerCamelCase__   	=					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():
            lowerCamelCase__   	=					hf_model(**lowerCamelCase    )

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

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

      lowerCamelCase__   	=					original_model.image_encoder(lowerCamelCase  ,					training=lowerCamelCase    )
      lowerCamelCase__   	=					original_model.text_encoder(lowerCamelCase  ,					training=lowerCamelCase    )

      lowerCamelCase__   	=					tf.nn.la_normalize(lowerCamelCase  ,					axis=-1    )
      lowerCamelCase__   	=					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__":
      __SCREAMING_SNAKE_CASE     : Optional[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""")

      __SCREAMING_SNAKE_CASE     : Any      =parser.parse_args()
      convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)





from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE     : str      =logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE     : Any      ={
    """ZinengTang/tvlt-base""": """https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json""",
}






class       A      (		_lowerCAmelCase		):
     """simple docstring"""

     _snake_case:						str  ="tvlt"




     def __init__(     self	: List[Any]						,	A__	: List[str]=2_2_4						,	A__	: Optional[int]=2_0_4_8						,	A__	: str=1_2_8						,	A__	: Union[str, Any]=[1_6, 1_6]						,	A__	: List[Any]=[1_6, 1_6]						,	A__	: int=3						,	A__	: Optional[int]=1						,	A__	: Tuple=8						,	A__	: Tuple=7_6_8						,	A__	: List[str]=1_2						,	A__	: Tuple=1_2						,	A__	: int=3_0_7_2						,	A__	: Tuple="gelu"						,	A__	: str=0.0						,	A__	: Dict=0.0						,	A__	: Any=0.02						,	A__	: int=1E-6						,	A__	: List[str]=True						,	A__	: List[str]=False						,	A__	: Optional[Any]=1_6						,	A__	: Optional[Any]=5_1_2						,	A__	: List[Any]=8						,	A__	: Any=2_0_4_8						,	A__	: List[Any]=0.75						,	A__	: List[Any]=0.15						,	A__	: str="frame-level"						,	A__	: Union[str, Any]=True						,	A__	: Optional[int]=True						,	A__	: Optional[Any]="classification"						,	**A__	: int						,	)					->							Tuple:
           super().__init__(**A__      )

           if audio_mask_type not in ("frame-level", "patch_level"):
                 raise ValueError(
                     'audio_mask_type must be one of two acceptable strategies - {\'frame_level\', \'patch-level\') '
                     F"""got {audio_mask_type}"""      )

           lowerCamelCase__   	=					image_size
           lowerCamelCase__   	=					spectrogram_length
           lowerCamelCase__   	=					frequency_length
           lowerCamelCase__   	=					image_patch_size
           lowerCamelCase__   	=					audio_patch_size
           lowerCamelCase__   	=					num_image_channels
           lowerCamelCase__   	=					num_audio_channels
           lowerCamelCase__   	=					num_frames

           lowerCamelCase__   	=					hidden_size
           lowerCamelCase__   	=					num_hidden_layers
           lowerCamelCase__   	=					num_attention_heads
           lowerCamelCase__   	=					intermediate_size
           lowerCamelCase__   	=					hidden_act
           lowerCamelCase__   	=					hidden_dropout_prob
           lowerCamelCase__   	=					attention_probs_dropout_prob
           lowerCamelCase__   	=					initializer_range
           lowerCamelCase__   	=					layer_norm_eps
           lowerCamelCase__   	=					qkv_bias
           lowerCamelCase__   	=					use_mean_pooling

           lowerCamelCase__   	=					decoder_num_attention_heads
           lowerCamelCase__   	=					decoder_hidden_size
           lowerCamelCase__   	=					decoder_num_hidden_layers
           lowerCamelCase__   	=					decoder_intermediate_size
           lowerCamelCase__   	=					pixel_mask_ratio
           lowerCamelCase__   	=					audio_mask_ratio
           lowerCamelCase__   	=					audio_mask_type

           lowerCamelCase__   	=					task_matching
           lowerCamelCase__   	=					task_mae
           lowerCamelCase__   	=					loss_type