



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()
UpperCAmelCase_          =				logging.get_logger(__name__)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Tuple					)->Tuple:
      _lowerCAmelCase			      =	tf.image.resize(_SCREAMING_SNAKE_CASE  ,					(3_4_6, 3_4_6)					)
      _lowerCAmelCase			      =	tf.image.crop_to_bounding_box(_SCREAMING_SNAKE_CASE  ,					(3_4_6 - 2_8_9) // 2  ,					(3_4_6 - 2_8_9) // 2  ,					2_8_9  ,					2_8_9					)
      return image
def 				UpperCAmelCase__		(					)->List[str]:
      _lowerCAmelCase			      =	EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''					)
      _lowerCAmelCase			      =	2_8_9
      _lowerCAmelCase			      =	6_4_0
      _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=_SCREAMING_SNAKE_CASE  ,					vision_config=_SCREAMING_SNAKE_CASE  ,					projection_dim=6_4_0					)
      return config
def 				UpperCAmelCase__		(					)->Optional[int]:
      _lowerCAmelCase			      =	'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      _lowerCAmelCase			      =	Image.open(requests.get(_SCREAMING_SNAKE_CASE  ,					stream=_SCREAMING_SNAKE_CASE					).raw					)
      return im
def 				UpperCAmelCase__		(					)->List[Any]:
      _lowerCAmelCase			      =	EfficientNetImageProcessor(
          do_center_crop=_SCREAMING_SNAKE_CASE  ,					rescale_factor=1 / 127.5  ,					rescale_offset=_SCREAMING_SNAKE_CASE  ,					do_normalize=_SCREAMING_SNAKE_CASE  ,					include_top=_SCREAMING_SNAKE_CASE  ,					resample=Image.BILINEAR  ,					)
      _lowerCAmelCase			      =	BertTokenizer.from_pretrained('''bert-base-uncased'''					)
      _lowerCAmelCase			      =	6_4
      _lowerCAmelCase			      =	AlignProcessor(image_processor=_SCREAMING_SNAKE_CASE  ,					tokenizer=_SCREAMING_SNAKE_CASE					)
      return processor
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Tuple					)->int:
      # EfficientNet image encoder
      _lowerCAmelCase			      =	[v.split('''_'''					)[0].split('''block'''					)[1] for v in original_param_names if v.startswith('''block'''					)]
      _lowerCAmelCase			      =	list(set(_SCREAMING_SNAKE_CASE					)					)
      _lowerCAmelCase			      =	sorted(_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	len(_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	{b: str(_SCREAMING_SNAKE_CASE					) for b, i in zip(_SCREAMING_SNAKE_CASE  ,					range(_SCREAMING_SNAKE_CASE					)					)}

      _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(1_2					):
            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 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						Optional[int]  ,					_SCREAMING_SNAKE_CASE :						List[Any]					)->Dict:
      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(_SCREAMING_SNAKE_CASE					).permute(3  ,					2  ,					0  ,					1					)
            elif "embeddings" in key:
                  _lowerCAmelCase			      =	torch.from_numpy(_SCREAMING_SNAKE_CASE					)
            elif "depthwise_kernel" in key:
                  _lowerCAmelCase			      =	torch.from_numpy(_SCREAMING_SNAKE_CASE					).permute(2  ,					3  ,					0  ,					1					)
            elif "kernel" in key:
                  _lowerCAmelCase			      =	torch.from_numpy(np.transpose(_SCREAMING_SNAKE_CASE					)					)
            elif "temperature" in key:
                  _lowerCAmelCase			      =	value
            elif "bn/gamma" or "bn/beta" in key:
                  _lowerCAmelCase			      =	torch.from_numpy(np.transpose(_SCREAMING_SNAKE_CASE					)					).squeeze()
            else:
                  _lowerCAmelCase			      =	torch.from_numpy(_SCREAMING_SNAKE_CASE					)

            # Replace HF parameters with original TF model parameters
            hf_params[hf_key].copy_(_SCREAMING_SNAKE_CASE					)
@torch.no_grad()
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						Tuple  ,					_SCREAMING_SNAKE_CASE :						Tuple  ,					_SCREAMING_SNAKE_CASE :						Dict					)->int:
      _lowerCAmelCase			      =	6_4
      _lowerCAmelCase			      =	Tokenizer(_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	align.Align('''efficientnet-b7'''  ,					'''bert-base'''  ,					6_4_0  ,					_SCREAMING_SNAKE_CASE  ,					tok.get_vocab_size()					)
      original_model.compile()
      original_model.load_weights(_SCREAMING_SNAKE_CASE					)

      _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(_SCREAMING_SNAKE_CASE					).eval()
      _lowerCAmelCase			      =	hf_model.state_dict()

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

      # Initialize processor
      _lowerCAmelCase			      =	get_processor()
      _lowerCAmelCase			      =	processor(
          images=prepare_img()  ,					text='''A picture of a cat'''  ,					padding='''max_length'''  ,					max_length=6_4  ,					return_tensors='''pt'''					)

      # HF model inference
      hf_model.eval()
      with torch.no_grad():
            _lowerCAmelCase			      =	hf_model(**_SCREAMING_SNAKE_CASE					)

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

      # Original model inference
      _lowerCAmelCase			      =	False
      _lowerCAmelCase			      =	EfficientNetImageProcessor(
          do_center_crop=_SCREAMING_SNAKE_CASE  ,					do_rescale=_SCREAMING_SNAKE_CASE  ,					do_normalize=_SCREAMING_SNAKE_CASE  ,					include_top=_SCREAMING_SNAKE_CASE  ,					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(_SCREAMING_SNAKE_CASE  ,					training=_SCREAMING_SNAKE_CASE					)
      _lowerCAmelCase			      =	original_model.text_encoder(_SCREAMING_SNAKE_CASE  ,					training=_SCREAMING_SNAKE_CASE					)

      _lowerCAmelCase			      =	tf.nn.la_normalize(_SCREAMING_SNAKE_CASE  ,					axis=-1					)
      _lowerCAmelCase			      =	tf.nn.la_normalize(_SCREAMING_SNAKE_CASE  ,					axis=-1					)

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

      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__":
    UpperCAmelCase_          =				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")

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




from typing import List, Union

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


UpperCAmelCase_          =				logging.get_logger(__name__)

UpperCAmelCase_          =				{
    "allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/config.json",
    # See all LED models at https://huggingface.co/models?filter=led
}





class 		UpperCAmelCase    (    snake_case_					):
   SCREAMING_SNAKE_CASE__					      =  '''led'''
   SCREAMING_SNAKE_CASE__					      =  {
       '''num_attention_heads''': '''encoder_attention_heads''',
       '''hidden_size''': '''d_model''',
       '''attention_probs_dropout_prob''': '''attention_dropout''',
       '''initializer_range''': '''init_std''',
   }





   def __init__(						self							,  _lowerCAmelCase=50_265							,  _lowerCAmelCase=16_384							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=12							,  _lowerCAmelCase=4_096							,  _lowerCAmelCase=16							,  _lowerCAmelCase=12							,  _lowerCAmelCase=4_096							,  _lowerCAmelCase=16							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=True							,  _lowerCAmelCase=True							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=0.1							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.02							,  _lowerCAmelCase=2							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=1							,  _lowerCAmelCase=0							,  _lowerCAmelCase=2							,  _lowerCAmelCase = 512							,  **_lowerCAmelCase							,  ):
         _lowerCAmelCase			      =	vocab_size
         _lowerCAmelCase			      =	max_encoder_position_embeddings
         _lowerCAmelCase			      =	max_decoder_position_embeddings
         _lowerCAmelCase			      =	d_model
         _lowerCAmelCase			      =	encoder_ffn_dim
         _lowerCAmelCase			      =	encoder_layers
         _lowerCAmelCase			      =	encoder_attention_heads
         _lowerCAmelCase			      =	decoder_ffn_dim
         _lowerCAmelCase			      =	decoder_layers
         _lowerCAmelCase			      =	decoder_attention_heads
         _lowerCAmelCase			      =	dropout
         _lowerCAmelCase			      =	attention_dropout
         _lowerCAmelCase			      =	activation_dropout
         _lowerCAmelCase			      =	activation_function
         _lowerCAmelCase			      =	init_std
         _lowerCAmelCase			      =	encoder_layerdrop
         _lowerCAmelCase			      =	decoder_layerdrop
         _lowerCAmelCase			      =	classifier_dropout
         _lowerCAmelCase			      =	use_cache
         _lowerCAmelCase			      =	encoder_layers
         _lowerCAmelCase			      =	attention_window

         super().__init__(
             pad_token_id=_lowerCAmelCase							,  bos_token_id=_lowerCAmelCase							,  eos_token_id=_lowerCAmelCase							,  is_encoder_decoder=_lowerCAmelCase							,  decoder_start_token_id=_lowerCAmelCase							,  **_lowerCAmelCase							,  )