

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__							=			logging.get_logger(__name__)
def   _lowerCAmelCase(						__A      ):
  UpperCAmelCase          =	tf.image.resize(__A ,							(346, 346)      )
  UpperCAmelCase          =	tf.image.crop_to_bounding_box(__A ,							(346 - 289) // 2 ,							(346 - 289) // 2 ,							289 ,							289      )
  return image
def   _lowerCAmelCase(						):
  UpperCAmelCase          =	EfficientNetConfig.from_pretrained("google/efficientnet-b7"      )
  UpperCAmelCase          =	289
  UpperCAmelCase          =	640
  UpperCAmelCase          =	{"0": "LABEL_0", "1": "LABEL_1"}
  UpperCAmelCase          =	{"LABEL_0": 0, "LABEL_1": 1}
  UpperCAmelCase          =	[]

  UpperCAmelCase          =	BertConfig()
  UpperCAmelCase          =	AlignConfig.from_text_vision_configs(
      text_config=__A ,							vision_config=__A ,							projection_dim=640      )
  return config
def   _lowerCAmelCase(						):
  UpperCAmelCase          =	"http://images.cocodataset.org/val2017/000000039769.jpg"
  UpperCAmelCase          =	Image.open(requests.get(__A ,							stream=__A      ).raw      )
  return im
def   _lowerCAmelCase(						):
  UpperCAmelCase          =	EfficientNetImageProcessor(
      do_center_crop=__A ,							rescale_factor=1 / 127.5 ,							rescale_offset=__A ,							do_normalize=__A ,							include_top=__A ,							resample=Image.BILINEAR ,							)
  UpperCAmelCase          =	BertTokenizer.from_pretrained("bert-base-uncased"      )
  UpperCAmelCase          =	64
  UpperCAmelCase          =	AlignProcessor(image_processor=__A ,							tokenizer=__A      )
  return processor
def   _lowerCAmelCase(						__A      ):
  # EfficientNet image encoder
  UpperCAmelCase          =	[v.split("_"      )[0].split("block"      )[1] for v in original_param_names if v.startswith("block"      )]
  UpperCAmelCase          =	list(set(__A      )      )
  UpperCAmelCase          =	sorted(__A      )
  UpperCAmelCase          =	len(__A      )
  UpperCAmelCase          =	{b: str(__A      ) for b, i in zip(__A ,							range(__A      )      )}

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

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

    # BERT text encoder
  UpperCAmelCase          =	[]
  UpperCAmelCase          =	"tf_bert_model/bert"
  UpperCAmelCase          =	"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:
      UpperCAmelCase          =	item[1]
  return key_mapping
def   _lowerCAmelCase(						__A ,							__A ,							__A      ):
  list(hf_params.keys()      )

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

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

    # Replace HF parameters with original TF model parameters
    hf_params[hf_key].copy_(__A      )
@torch.no_grad()
def   _lowerCAmelCase(						__A ,							__A ,							__A ,							__A      ):
  UpperCAmelCase          =	64
  UpperCAmelCase          =	Tokenizer(__A      )
  UpperCAmelCase          =	align.Align("efficientnet-b7" ,							"bert-base" ,							640 ,							__A ,							tok.get_vocab_size()      )
  original_model.compile()
  original_model.load_weights(__A      )

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

  # Load HuggingFace model
  UpperCAmelCase          =	get_align_config()
  UpperCAmelCase          =	AlignModel(__A      ).eval()
  UpperCAmelCase          =	hf_model.state_dict()

  # Create src-to-dst parameter name mapping dictionary
  print("Converting parameters..."      )
  UpperCAmelCase          =	rename_keys(__A      )
  replace_params(__A ,							__A ,							__A      )

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

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

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

  UpperCAmelCase          =	original_model.image_encoder(__A ,							training=__A      )
  UpperCAmelCase          =	original_model.text_encoder(__A ,							training=__A      )

  UpperCAmelCase          =	tf.nn.la_normalize(__A ,							axis=-1      )
  UpperCAmelCase          =	tf.nn.la_normalize(__A ,							axis=-1      )

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

  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__							=			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__							=			parser.parse_args()
      convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)






import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available
def   _lowerCAmelCase(						__A ,							__A ,							__A ,							__A=None      ):
  UpperCAmelCase          =	None
  # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
  if str(weight.dtype      ) == "torch.bfloat16":
    # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
    UpperCAmelCase          =	weight.view(torch.intaa      )
    UpperCAmelCase          =	"bfloat16"
  UpperCAmelCase          =	weight.cpu().numpy()
  UpperCAmelCase          =	os.path.join(__A ,							F"{weight_name}.dat"      )
  if index is not None:
    if dtype is None:
      UpperCAmelCase          =	str(array.dtype      )
    UpperCAmelCase          =	{"dtype": dtype, "shape": list(array.shape      )}
  if array.ndim == 0:
    UpperCAmelCase          =	array[None]
  UpperCAmelCase          =	np.memmap(__A ,							dtype=array.dtype ,							mode="w+" ,							shape=array.shape      )
  UpperCAmelCase          =	array[:]
  file_array.flush()
  return index
def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	tuple(weight_info["shape"]      )
  if shape == ():
    # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
    UpperCAmelCase          =	(1,)

  UpperCAmelCase          =	weight_info["dtype"]
  if dtype == "bfloat16":
    # NumPy does not support bfloat16 so this was saved as a int16
    UpperCAmelCase          =	"int16"

  UpperCAmelCase          =	np.memmap(__A ,							dtype=__A ,							shape=__A ,							mode="r"      )

  if len(weight_info["shape"]      ) == 0:
    UpperCAmelCase          =	weight[0]
  UpperCAmelCase          =	torch.tensor(__A      )
  if weight_info["dtype"] == "bfloat16":
    UpperCAmelCase          =	weight.view(torch.bfloataa      )

  return weight
def   _lowerCAmelCase(						__A ,							__A      ):
  if index is None or len(__A      ) == 0:
    # Nothing to save
    return

  UpperCAmelCase          =	os.path.join(__A ,							"index.json"      )
  if os.path.isfile(__A      ):
    with open(__A ,							"r" ,							encoding="utf-8"      ) as f:
      UpperCAmelCase          =	json.load(__A      )
  else:
    UpperCAmelCase          =	{}
  current_index.update(__A      )

  with open(__A ,							"w" ,							encoding="utf-8"      ) as f:
    json.dump(__A ,							__A ,							indent=2      )
def   _lowerCAmelCase(						__A ,							__A      ):
  os.makedirs(__A ,							exist_ok=__A      )
  UpperCAmelCase          =	{}
  for name, parameter in state_dict.items():
    UpperCAmelCase          =	offload_weight(__A ,							__A ,							__A ,							index=__A      )

  # Update index
  save_offload_index(__A ,							__A      )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Tuple      ,    lowerCAmelCase__		:			Mapping      ,    lowerCAmelCase__		:			str     )     ->   Tuple:
   UpperCAmelCase          =	dataset
   UpperCAmelCase          =	prefix


 def __getitem__(    self		:			Dict      ,    lowerCAmelCase__		:			Any     )     ->   Optional[int]:
   return self.dataset[f"{self.prefix}{key}"]


 def __iter__(    self		:			Optional[int]     )     ->   Any:
   return iter([key for key in self.dataset if key.startswith(self.prefix     )]     )


 def __len__(    self		:			Union[str, Any]     )     ->   int:
   return len(self.dataset     )






class     __magic_name__     (							_snake_case     ):


 def __init__(    self		:			Optional[int]      ,    lowerCAmelCase__		:			Dict[str, torch.Tensor] = None      ,    lowerCAmelCase__		:			Optional[Union[str, os.PathLike]] = None      ,    lowerCAmelCase__		:			Mapping = None      ,    lowerCAmelCase__		:			Dict=None      ,    )     ->   Dict:
   if state_dict is None and save_folder is None:
     raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights."     )

   UpperCAmelCase          =	{} if state_dict is None else state_dict
   UpperCAmelCase          =	save_folder
   if index is None and save_folder is not None:
     with open(os.path.join(lowerCAmelCase__      ,    "index.json"     )     ) as f:
       UpperCAmelCase          =	json.load(lowerCAmelCase__     )
   UpperCAmelCase          =	{} if index is None else index
   UpperCAmelCase          =	list(self.state_dict.keys()     )
   self.all_keys.extend([key for key in self.index if key not in self.all_keys]     )
   UpperCAmelCase          =	device


 def __getitem__(    self		:			str      ,    lowerCAmelCase__		:			str     )     ->   Dict:
   # State dict gets priority
   if key in self.state_dict:
     return self.state_dict[key]
   UpperCAmelCase          =	self.index[key]
   if weight_info.get("safetensors_file"     ) is not None:
     if not is_safetensors_available():
       raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`."     )

     from safetensors import safe_open

     UpperCAmelCase          =	"cpu" if self.device is None else self.device
     with safe_open(weight_info["safetensors_file"]      ,    framework="pt"      ,    device=lowerCAmelCase__     ) as f:
       UpperCAmelCase          =	f.get_tensor(weight_info.get("weight_name"      ,    lowerCAmelCase__     )     )

     if "dtype" in weight_info:
       return tensor.to(getattr(lowerCAmelCase__      ,    weight_info["dtype"]     )     )
     else:
       return tensor

   UpperCAmelCase          =	os.path.join(self.save_folder      ,    f"{key}.dat"     )
   return load_offloaded_weight(lowerCAmelCase__      ,    lowerCAmelCase__     )


 def __iter__(    self		:			str     )     ->   Dict:
   return iter(self.all_keys     )


 def __len__(    self		:			Optional[Any]     )     ->   List[Any]:
   return len(self.all_keys     )

def   _lowerCAmelCase(						__A ,							__A      ):
  UpperCAmelCase          =	{}
  for module_name in submodule_names:
    # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
    # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
    result.update(
        {
            key: param
            for key, param in state_dict.items()
            if key == module_name or key.startswith(module_name + "."      )
        }      )
  return result



