


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






def  _UpperCAmelCase						( __lowerCamelCase							)		->List[str]:


 '''simple docstring'''
 snake_case :			Optional[Any]							=   tf.image.resize(_UpperCAmelCase			,  (346, 346)							)
 snake_case :			Optional[int]							=   tf.image.crop_to_bounding_box(_UpperCAmelCase			,  (346 - 289) // 2			,  (346 - 289) // 2			,  289			,  289							)
 return image






def  _UpperCAmelCase						( )		->Optional[int]:


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

 snake_case :			str							=   BertConfig()
 snake_case :			str							=   AlignConfig.from_text_vision_configs(
     text_config=_UpperCAmelCase			,  vision_config=_UpperCAmelCase			,  projection_dim=640							)
 return config






def  _UpperCAmelCase						( )		->List[Any]:


 '''simple docstring'''
 snake_case :			Dict							=   '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 snake_case :			List[str]							=   Image.open(requests.get(_UpperCAmelCase			,  stream=_UpperCAmelCase							).raw							)
 return im






def  _UpperCAmelCase						( )		->List[Any]:


 '''simple docstring'''
 snake_case :			Dict							=   EfficientNetImageProcessor(
     do_center_crop=_UpperCAmelCase			,  rescale_factor=1 / 127.5			,  rescale_offset=_UpperCAmelCase			,  do_normalize=_UpperCAmelCase			,  include_top=_UpperCAmelCase			,  resample=Image.BILINEAR			,  )
 snake_case :			Tuple							=   BertTokenizer.from_pretrained('''bert-base-uncased'''							)
 snake_case :			int							=   64
 snake_case :			str							=   AlignProcessor(image_processor=_UpperCAmelCase			,  tokenizer=_UpperCAmelCase							)
 return processor






def  _UpperCAmelCase						( __lowerCamelCase							)		->List[str]:


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

 snake_case :			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:
  snake_case :			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''')							)

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

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






def  _UpperCAmelCase						( __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase							)		->List[Any]:


 '''simple docstring'''
 list(hf_params.keys()							)

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

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

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






@torch.no_grad()
def  _UpperCAmelCase						( __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase							)		->List[Any]:


 '''simple docstring'''
 snake_case :			List[Any]							=   64
 snake_case :			Optional[int]							=   Tokenizer(_UpperCAmelCase							)
 snake_case :			List[Any]							=   align.Align('''efficientnet-b7'''			,  '''bert-base'''			,  640			,  _UpperCAmelCase			,  tok.get_vocab_size()							)
 original_model.compile()
 original_model.load_weights(_UpperCAmelCase							)

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

 # Load HuggingFace model
 snake_case :			List[str]							=   get_align_config()
 snake_case :			Optional[Any]							=   AlignModel(_UpperCAmelCase							).eval()
 snake_case :			List[str]							=   hf_model.state_dict()

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

 # Initialize processor
 snake_case :			Optional[Any]							=   get_processor()
 snake_case :			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():
  snake_case :			Tuple							=   hf_model(**_UpperCAmelCase							)

 snake_case :			Tuple							=   outputs.image_embeds.detach().numpy()
 snake_case :			Any							=   outputs.text_embeds.detach().numpy()

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

 snake_case :			List[Any]							=   original_model.image_encoder(_UpperCAmelCase			,  training=_UpperCAmelCase							)
 snake_case :			str							=   original_model.text_encoder(_UpperCAmelCase			,  training=_UpperCAmelCase							)

 snake_case :			List[Any]							=   tf.nn.la_normalize(_UpperCAmelCase			,  axis=-1							)
 snake_case :			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					:				Optional[int]       =      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					:				str       =      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_torch_available,
)


UpperCAmelCase__  =			{
    "configuration_align": [
        "ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "AlignConfig",
        "AlignTextConfig",
        "AlignVisionConfig",
    ],
    "processing_align": ["AlignProcessor"],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__  =			[
							    "ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "AlignModel",
							    "AlignPreTrainedModel",
							    "AlignTextModel",
							    "AlignVisionModel",
							]

if TYPE_CHECKING:
							from .configuration_align import (
							    ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    AlignConfig,
							    AlignTextConfig,
							    AlignVisionConfig,
							)
							from .processing_align import AlignProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_align import (
														    ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST,
														    AlignModel,
														    AlignPreTrainedModel,
														    AlignTextModel,
														    AlignVisionModel,
														)

else:
							import sys

							UpperCAmelCase__  =			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
