





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

def      __snake_case						(					_lowerCAmelCase : Optional[Any]	) ->						Any:
   A_					:				int          =			tf.image.resize(_lowerCAmelCase				,				(346, 346)	)
   A_					:				Dict          =			tf.image.crop_to_bounding_box(_lowerCAmelCase				,				(346 - 289) // 2				,				(346 - 289) // 2				,				289				,				289	)
   return image

def      __snake_case						(					) ->						Tuple:
   A_					:				str          =			EfficientNetConfig.from_pretrained("google/efficientnet-b7"	)
   A_					:				Optional[int]          =			289
   A_					:				Union[str, Any]          =			640
   A_					:				Dict          =			{"0": "LABEL_0", "1": "LABEL_1"}
   A_					:				List[Any]          =			{"LABEL_0": 0, "LABEL_1": 1}
   A_					:				List[str]          =			[]

   A_					:				List[Any]          =			BertConfig()
   A_					:				Any          =			AlignConfig.from_text_vision_configs(
       text_config=_lowerCAmelCase				,				vision_config=_lowerCAmelCase				,				projection_dim=640	)
   return config

def      __snake_case						(					) ->						List[str]:
   A_					:				str          =			"http://images.cocodataset.org/val2017/000000039769.jpg"
   A_					:				Optional[Any]          =			Image.open(requests.get(_lowerCAmelCase				,				stream=_lowerCAmelCase	).raw	)
   return im

def      __snake_case						(					) ->						Optional[int]:
   A_					:				List[str]          =			EfficientNetImageProcessor(
       do_center_crop=_lowerCAmelCase				,				rescale_factor=1 / 1_27.5				,				rescale_offset=_lowerCAmelCase				,				do_normalize=_lowerCAmelCase				,				include_top=_lowerCAmelCase				,				resample=Image.BILINEAR				,				)
   A_					:				Optional[Any]          =			BertTokenizer.from_pretrained("bert-base-uncased"	)
   A_					:				Tuple          =			64
   A_					:				List[Any]          =			AlignProcessor(image_processor=_lowerCAmelCase				,				tokenizer=_lowerCAmelCase	)
   return processor

def      __snake_case						(					_lowerCAmelCase : str	) ->						Optional[Any]:
   # EfficientNet image encoder
   A_					:				Tuple          =			[v.split("_"	)[0].split("block"	)[1] for v in original_param_names if v.startswith("block"	)]
   A_					:				Union[str, Any]          =			list(set(_lowerCAmelCase	)	)
   A_					:				Union[str, Any]          =			sorted(_lowerCAmelCase	)
   A_					:				Any          =			len(_lowerCAmelCase	)
   A_					:				str          =			{b: str(_lowerCAmelCase	) for b, i in zip(_lowerCAmelCase				,				range(_lowerCAmelCase	)	)}

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

   A_					:				List[str]          =			{}
   for item in rename_keys:
      if item[0] in original_param_names:
         A_					:				Optional[int]          =			"vision_model." + item[1]

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

def      __snake_case						(					_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : Tuple				,				_lowerCAmelCase : Tuple	) ->						Union[str, Any]:
   list(hf_params.keys()	)

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

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

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

@torch.no_grad()
def      __snake_case						(					_lowerCAmelCase : List[Any]				,				_lowerCAmelCase : Any				,				_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : int	) ->						List[Any]:
   A_					:				Dict          =			64
   A_					:				Dict          =			Tokenizer(_lowerCAmelCase	)
   A_					:				Union[str, Any]          =			align.Align("efficientnet-b7"				,				"bert-base"				,				640				,				_lowerCAmelCase				,				tok.get_vocab_size()	)
   original_model.compile()
   original_model.load_weights(_lowerCAmelCase	)

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

   # Load HuggingFace model
   A_					:				Optional[int]          =			get_align_config()
   A_					:				List[Any]          =			AlignModel(_lowerCAmelCase	).eval()
   A_					:				Dict          =			hf_model.state_dict()

   # Create src-to-dst parameter name mapping dictionary
   print("Converting parameters..."	)
   A_					:				Dict          =			rename_keys(_lowerCAmelCase	)
   replace_params(_lowerCAmelCase				,				_lowerCAmelCase				,				_lowerCAmelCase	)

   # Initialize processor
   A_					:				int          =			get_processor()
   A_					:				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():
      A_					:				Optional[Any]          =			hf_model(**_lowerCAmelCase	)

   A_					:				List[str]          =			outputs.image_embeds.detach().numpy()
   A_					:				Tuple          =			outputs.text_embeds.detach().numpy()

   # Original model inference
   A_					:				List[str]          =			False
   A_					:				List[str]          =			EfficientNetImageProcessor(
       do_center_crop=_lowerCAmelCase				,				do_rescale=_lowerCAmelCase				,				do_normalize=_lowerCAmelCase				,				include_top=_lowerCAmelCase				,				resample=Image.BILINEAR				,				)
   A_					:				List[str]          =			tf_image_processor(images=prepare_img()				,				return_tensors="tf"				,				data_format="channels_last"	)["pixel_values"]
   A_					:				Dict          =			tok(tf.constant(["A picture of a cat"]	)	)

   A_					:				List[Any]          =			original_model.image_encoder(_lowerCAmelCase				,				training=_lowerCAmelCase	)
   A_					:				List[Any]          =			original_model.text_encoder(_lowerCAmelCase				,				training=_lowerCAmelCase	)

   A_					:				Any          =			tf.nn.la_normalize(_lowerCAmelCase				,				axis=-1	)
   A_					:				Tuple          =			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__":
 _lowerCAmelCase							:      List[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''')

 _lowerCAmelCase							:      List[str]						=	parser.parse_args()
 convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)







import argparse

import torch

from transformers import ChineseCLIPConfig, ChineseCLIPModel

def      __snake_case						(					_lowerCAmelCase : str				,				_lowerCAmelCase : Optional[int]				,				_lowerCAmelCase : Optional[Any]	) ->						List[str]:
   A_      ,						A_      ,						A_					:				Dict          =			pt_weights[f"{prefix}.in_proj_weight"].chunk(3				,				dim=0	)
   A_      ,						A_      ,						A_					:				Union[str, Any]          =			pt_weights[f"{prefix}.in_proj_bias"].chunk(3				,				dim=0	)

   A_					:				int          =			pt_weights[f"{prefix}.out_proj.weight"]
   A_					:				Dict          =			pt_weights[f"{prefix}.out_proj.bias"]

   A_					:				Optional[int]          =			q_proj
   A_					:				Optional[int]          =			q_proj_bias

   A_					:				Tuple          =			k_proj
   A_					:				str          =			k_proj_bias

   A_					:				Optional[int]          =			v_proj
   A_					:				int          =			v_proj_bias

   A_					:				int          =			out_proj_weights
   A_					:				List[Any]          =			out_proj_bias

def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : Dict				,				_lowerCAmelCase : int	) ->						Dict:
   copy_linear(hf_mlp.fca				,				_lowerCAmelCase				,				f"{prefix}.c_fc"	)
   copy_linear(hf_mlp.fca				,				_lowerCAmelCase				,				f"{prefix}.c_proj"	)

def      __snake_case						(					_lowerCAmelCase : Tuple				,				_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : Optional[int]	) ->						List[str]:
   A_					:				List[str]          =			pt_weights[f"{prefix}.weight"].data
   A_					:				int          =			pt_weights[f"{prefix}.bias"].data

def      __snake_case						(					_lowerCAmelCase : Dict				,				_lowerCAmelCase : str				,				_lowerCAmelCase : Optional[Any]	) ->						str:
   # copy layer norms
   copy_linear(hf_layer.layer_norma				,				_lowerCAmelCase				,				f"{prefix}.ln_1"	)
   copy_linear(hf_layer.layer_norma				,				_lowerCAmelCase				,				f"{prefix}.ln_2"	)

   # copy MLP
   copy_mlp(hf_layer.mlp				,				_lowerCAmelCase				,				f"{prefix}.mlp"	)

   # copy attn
   copy_attn_layer(hf_layer.self_attn				,				_lowerCAmelCase				,				f"{prefix}.attn"	)

def      __snake_case						(					_lowerCAmelCase : Optional[Any]				,				_lowerCAmelCase : int				,				_lowerCAmelCase : Optional[Any]	) ->						List[Any]:
   for layer_id, hf_layer in enumerate(_lowerCAmelCase	):
      copy_layer(_lowerCAmelCase				,				_lowerCAmelCase				,				f"{prefix}.{layer_id}"	)

def      __snake_case						(					_lowerCAmelCase : Union[str, Any]				,				_lowerCAmelCase : Any	) ->						Dict:
   # copy projection
   A_					:				Optional[Any]          =			pt_weights["text_projection"].data.T

   # copy text encoder
   for name, param in hf_model.text_model.named_parameters():
      A_					:				Any          =			pt_weights[f"bert.{name}"].data

def      __snake_case						(					_lowerCAmelCase : int				,				_lowerCAmelCase : Any	) ->						List[str]:
   # copy projection
   A_					:				Tuple          =			pt_weights["visual.proj"].data.T

   # copy layer norms
   copy_linear(hf_model.vision_model.pre_layrnorm				,				_lowerCAmelCase				,				"visual.ln_pre"	)
   copy_linear(hf_model.vision_model.post_layernorm				,				_lowerCAmelCase				,				"visual.ln_post"	)

   # copy embeddings
   A_					:				List[Any]          =			pt_weights["visual.conv1.weight"].data
   A_					:				Union[str, Any]          =			pt_weights["visual.class_embedding"].data
   A_					:				Tuple          =			pt_weights["visual.positional_embedding"].data

   # copy encoder
   copy_layers(hf_model.vision_model.encoder.layers				,				_lowerCAmelCase				,				"visual.transformer.resblocks"	)

@torch.no_grad()
def      __snake_case						(					_lowerCAmelCase : List[Any]				,				_lowerCAmelCase : int				,				_lowerCAmelCase : Optional[int]=None	) ->						str:
   assert config_path is not None, "Please specify the ChineseCLIP model config of the corresponding model size."
   A_					:				Optional[Any]          =			ChineseCLIPConfig.from_pretrained(_lowerCAmelCase	)

   A_					:				Tuple          =			ChineseCLIPModel(_lowerCAmelCase	).eval()

   A_					:				Any          =			torch.load(_lowerCAmelCase				,				map_location="cpu"	)["state_dict"]
   A_					:				Optional[Any]          =			{(name[7:] if name.startswith("module."	) else name): value for name, value in pt_weights.items()}

   copy_text_model_and_projection(_lowerCAmelCase				,				_lowerCAmelCase	)
   copy_vision_model_and_projection(_lowerCAmelCase				,				_lowerCAmelCase	)
   A_					:				List[Any]          =			pt_weights["logit_scale"].data

   hf_model.save_pretrained(_lowerCAmelCase	)


if __name__ == "__main__":
 _lowerCAmelCase							:      Optional[Any]						=	argparse.ArgumentParser()
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     help='''Path to the output folder storing converted hf PyTorch model.''',
 )
 parser.add_argument(
     '''--checkpoint_path''', default=None, type=str, help='''Path to original github format ChineseCLIP checkpoint.'''
 )
 parser.add_argument(
     '''--config_path''', default=None, required=True, type=str, help='''Path to hf config.json of model to convert.'''
 )
 _lowerCAmelCase							:      Tuple						=	parser.parse_args()

 convert_chinese_clip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
 print('''The conversion is finished!''')
