codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert align checkpoints from the original repository 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 setverbosityinfo logger logging getloggername def preprocessimage image tf image resizeimage 346 346 image tf image croptoboundingboximage 346 289 2 346 289 2 289 289 return image def getalignconfig visionconfig efficientnetconfig frompretrainedgoogleefficientnetb7 visionconfig imagesize 289 visionconfig hiddendim 640 visionconfig id2label 0 label0 1 label1 visionconfig label2id label0 0 label1 1 visionconfig depthwisepadding textconfig bertconfig config alignconfig fromtextvisionconfigs textconfigtextconfig visionconfigvisionconfig projectiondim640 return config we will verify our results on an image of cute cats def prepareimg url http images cocodataset orgval2017000000039769 jpg im image openrequests geturl streamtrue raw return im def getprocessor imageprocessor efficientnetimageprocessor docentercroptrue rescalefactor1 127 5 rescaleoffsettrue donormalizefalse includetopfalse resampleimage bilinear tokenizer berttokenizer frompretrainedbertbaseuncased tokenizer modelmaxlength 64 processor alignprocessorimageprocessorimageprocessor tokenizertokenizer return processor here we list all keys to be renamed original name on the left our name on the right def renamekeysoriginalparamnames efficientnet image encoder blocknames v split0 splitblock1 for v in originalparamnames if v startswithblock blocknames listsetblocknames blocknames sortedblocknames numblocks lenblocknames blocknamemapping b stri for b i in zipblocknames rangenumblocks renamekeys renamekeys appendstemconvkernel 0 embeddings convolution weight renamekeys appendstembngamma 0 embeddings batchnorm weight renamekeys appendstembnbeta 0 embeddings batchnorm bias renamekeys appendstembnmovingmean 0 embeddings batchnorm runningmean renamekeys appendstembnmovingvariance 0 embeddings batchnorm runningvar for b in blocknames hfb blocknamemappingb renamekeys appendfblockbexpandconvkernel 0 fencoder blocks hfb expansion expandconv weight renamekeys appendfblockbexpandbngamma 0 fencoder blocks hfb expansion expandbn weight renamekeys appendfblockbexpandbnbeta 0 fencoder blocks hfb expansion expandbn bias renamekeys append fblockbexpandbnmovingmean 0 fencoder blocks hfb expansion expandbn runningmean renamekeys append fblockbexpandbnmovingvariance 0 fencoder blocks hfb expansion expandbn runningvar renamekeys append fblockbdwconvdepthwisekernel 0 fencoder blocks hfb depthwiseconv depthwiseconv weight renamekeys appendfblockbbngamma 0 fencoder blocks hfb depthwiseconv depthwisenorm weight renamekeys appendfblockbbnbeta 0 fencoder blocks hfb depthwiseconv depthwisenorm bias renamekeys append fblockbbnmovingmean 0 fencoder blocks hfb depthwiseconv depthwisenorm runningmean renamekeys append fblockbbnmovingvariance 0 fencoder blocks hfb depthwiseconv depthwisenorm runningvar renamekeys appendfblockbsereducekernel 0 fencoder blocks hfb squeezeexcite reduce weight renamekeys appendfblockbsereducebias 0 fencoder blocks hfb squeezeexcite reduce bias renamekeys appendfblockbseexpandkernel 0 fencoder blocks hfb squeezeexcite expand weight renamekeys appendfblockbseexpandbias 0 fencoder blocks hfb squeezeexcite expand bias renamekeys append fblockbprojectconvkernel 0 fencoder blocks hfb projection projectconv weight renamekeys appendfblockbprojectbngamma 0 fencoder blocks hfb projection projectbn weight renamekeys appendfblockbprojectbnbeta 0 fencoder blocks hfb projection projectbn bias renamekeys append fblockbprojectbnmovingmean 0 fencoder blocks hfb projection projectbn runningmean renamekeys append fblockbprojectbnmovingvariance 0 fencoder blocks hfb projection projectbn runningvar keymapping for item in renamekeys if item0 in originalparamnames keymappingitem0 visionmodel item1 bert text encoder renamekeys old tfbertmodelbert new textmodel for i in range12 renamekeys append foldencoderlayer iattentionselfquerykernel 0 fnew encoder layer i attention self query weight renamekeys append foldencoderlayer iattentionselfquerybias 0 fnew encoder layer i attention self query bias renamekeys append foldencoderlayer iattentionselfkeykernel 0 fnew encoder layer i attention self key weight renamekeys append foldencoderlayer iattentionselfkeybias 0 fnew encoder layer i attention self key bias renamekeys append foldencoderlayer iattentionselfvaluekernel 0 fnew encoder layer i attention self value weight renamekeys append foldencoderlayer iattentionselfvaluebias 0 fnew encoder layer i attention self value bias renamekeys append foldencoderlayer iattentionoutputdensekernel 0 fnew encoder layer i attention output dense weight renamekeys append foldencoderlayer iattentionoutputdensebias 0 fnew encoder layer i attention output dense bias renamekeys append foldencoderlayer iattentionoutputlayernormgamma 0 fnew encoder layer i attention output layernorm weight renamekeys append foldencoderlayer iattentionoutputlayernormbeta 0 fnew encoder layer i attention output layernorm bias renamekeys append foldencoderlayer iintermediatedensekernel 0 fnew encoder layer i intermediate dense weight renamekeys append foldencoderlayer iintermediatedensebias 0 fnew encoder layer i intermediate dense bias renamekeys append foldencoderlayer ioutputdensekernel 0 fnew encoder layer i output dense weight renamekeys append foldencoderlayer ioutputdensebias 0 fnew encoder layer i output dense bias renamekeys append foldencoderlayer ioutputlayernormgamma 0 fnew encoder layer i output layernorm weight renamekeys append foldencoderlayer ioutputlayernormbeta 0 fnew encoder layer i output layernorm bias renamekeys appendfoldembeddingswordembeddingsweight 0 fnew embeddings wordembeddings weight renamekeys append foldembeddingspositionembeddingsembeddings 0 fnew embeddings positionembeddings weight renamekeys append foldembeddingstokentypeembeddingsembeddings 0 fnew embeddings tokentypeembeddings weight renamekeys appendfoldembeddingslayernormgamma 0 fnew embeddings layernorm weight renamekeys appendfoldembeddingslayernormbeta 0 fnew embeddings layernorm bias renamekeys appendfoldpoolerdensekernel 0 fnew pooler dense weight renamekeys appendfoldpoolerdensebias 0 fnew pooler dense bias renamekeys appenddensekernel 0 textprojection weight renamekeys appenddensebias 0 textprojection bias renamekeys appenddensebias 0 textprojection bias renamekeys appendtemperature 0 temperature for item in renamekeys if item0 in originalparamnames keymappingitem0 item1 return keymapping def replaceparamshfparams tfparams keymapping listhfparams keys for key value in tfparams items if key not in keymapping continue hfkey keymappingkey if conv in key and kernel in key newhfvalue torch fromnumpyvalue permute3 2 0 1 elif embeddings in key newhfvalue torch fromnumpyvalue elif depthwisekernel in key newhfvalue torch fromnumpyvalue permute2 3 0 1 elif kernel in key newhfvalue torch fromnumpynp transposevalue elif temperature in key newhfvalue value elif bngamma or bnbeta in key newhfvalue torch fromnumpynp transposevalue squeeze else newhfvalue torch fromnumpyvalue replace hf parameters with original tf model parameters hfparamshfkey copynewhfvalue torch nograd def convertaligncheckpointcheckpointpath pytorchdumpfolderpath savemodel pushtohub load original model seqlength 64 tok tokenizerseqlength originalmodel align alignefficientnetb7 bertbase 640 seqlength tok getvocabsize originalmodel compile originalmodel loadweightscheckpointpath tfparams originalmodel trainablevariables tfnontrainparams originalmodel nontrainablevariables tfparams param name param numpy for param in tfparams for param in tfnontrainparams tfparamsparam name param numpy tfparamnames listtfparams keys load huggingface model config getalignconfig hfmodel alignmodelconfig eval hfparams hfmodel statedict create srctodst parameter name mapping dictionary printconverting parameters keymapping renamekeystfparamnames replaceparamshfparams tfparams keymapping initialize processor processor getprocessor inputs processor imagesprepareimg texta picture of a cat paddingmaxlength maxlength64 returntensorspt hf model inference hfmodel eval with torch nograd outputs hfmodelinputs hfimagefeatures outputs imageembeds detach numpy hftextfeatures outputs textembeds detach numpy original model inference originalmodel trainable false tfimageprocessor efficientnetimageprocessor docentercroptrue dorescalefalse donormalizefalse includetopfalse resampleimage bilinear image tfimageprocessorimagesprepareimg returntensorstf dataformatchannelslastpixelvalues text toktf constanta picture of a cat imagefeatures originalmodel imageencoderimage trainingfalse textfeatures originalmodel textencodertext trainingfalse imagefeatures tf nn l2normalizeimagefeatures axis1 textfeatures tf nn l2normalizetextfeatures axis1 check whether original and hf model outputs match np allclose if not np allcloseimagefeatures hfimagefeatures atol1e3 raise valueerrorthe predicted image features are not the same if not np allclosetextfeatures hftextfeatures atol1e3 raise valueerrorthe predicted text features are not the same printmodel outputs match if savemodel create folder to save model if not os path isdirpytorchdumpfolderpath os mkdirpytorchdumpfolderpath save converted model and image processor hfmodel savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub push model and image processor to hub printpushing converted align to the hub processor pushtohubalignbase hfmodel pushtohubalignbase if name main parser argparse argumentparser required parameters parser addargument checkpointpath default weightsmodelweights typestr helppath to the pretrained tf align checkpoint parser addargument pytorchdumpfolderpath defaulthfmodel typestr helppath to the output pytorch model directory parser addargumentsavemodel actionstoretrue helpsave model to local parser addargumentpushtohub actionstoretrue helppush model and image processor to the hub args parser parseargs convertaligncheckpointargs checkpointpath args pytorchdumpfolderpath args savemodel args pushtohub coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert align checkpoints from the original repository we will verify our results on an image of cute cats here we list all keys to be renamed original name on the left our name on the right efficientnet image encoder bert text encoder replace hf parameters with original tf model parameters copy paste tweak model s weights to our align structure load original model load huggingface model create src to dst parameter name mapping dictionary initialize processor hf model inference original model inference check whether original and hf model outputs match np allclose create folder to save model save converted model and image processor push model and image processor to hub required parameters
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()
logger = logging.get_logger(__name__)


def preprocess(image):
    image = tf.image.resize(image, (346, 346))
    image = tf.image.crop_to_bounding_box(image, (346 - 289) // 2, (346 - 289) // 2, 289, 289)
    return image


def get_align_config():
    vision_config = EfficientNetConfig.from_pretrained("google/efficientnet-b7")
    vision_config.image_size = 289
    vision_config.hidden_dim = 640
    vision_config.id2label = {"0": "LABEL_0", "1": "LABEL_1"}
    vision_config.label2id = {"LABEL_0": 0, "LABEL_1": 1}
    vision_config.depthwise_padding = []

    text_config = BertConfig()
    config = AlignConfig.from_text_vision_configs(
        text_config=text_config, vision_config=vision_config, projection_dim=640
    )
    return config



def prepare_img():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    im = Image.open(requests.get(url, stream=True).raw)
    return im


def get_processor():
    image_processor = EfficientNetImageProcessor(
        do_center_crop=True,
        rescale_factor=1 / 127.5,
        rescale_offset=True,
        do_normalize=False,
        include_top=False,
        resample=Image.BILINEAR,
    )
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    tokenizer.model_max_length = 64
    processor = AlignProcessor(image_processor=image_processor, tokenizer=tokenizer)
    return processor



def rename_keys(original_param_names):
    
    block_names = [v.split("_")[0].split("block")[1] for v in original_param_names if v.startswith("block")]
    block_names = list(set(block_names))
    block_names = sorted(block_names)
    num_blocks = len(block_names)
    block_name_mapping = {b: str(i) for b, i in zip(block_names, range(num_blocks))}

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

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

    
    rename_keys = []
    old = "tf_bert_model/bert"
    new = "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:
            key_mapping[item[0]] = item[1]
    return key_mapping


def replace_params(hf_params, tf_params, key_mapping):
    list(hf_params.keys())

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

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

        
        hf_params[hf_key].copy_(new_hf_value)


@torch.no_grad()
def convert_align_checkpoint(checkpoint_path, pytorch_dump_folder_path, save_model, push_to_hub):
    
    
    seq_length = 64
    tok = Tokenizer(seq_length)
    original_model = align.Align("efficientnet-b7", "bert-base", 640, seq_length, tok.get_vocab_size())
    original_model.compile()
    original_model.load_weights(checkpoint_path)

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

    
    config = get_align_config()
    hf_model = AlignModel(config).eval()
    hf_params = hf_model.state_dict()

    
    print("Converting parameters...")
    key_mapping = rename_keys(tf_param_names)
    replace_params(hf_params, tf_params, key_mapping)

    
    processor = get_processor()
    inputs = processor(
        images=prepare_img(), text="A picture of a cat", padding="max_length", max_length=64, return_tensors="pt"
    )

    
    hf_model.eval()
    with torch.no_grad():
        outputs = hf_model(**inputs)

    hf_image_features = outputs.image_embeds.detach().numpy()
    hf_text_features = outputs.text_embeds.detach().numpy()

    
    original_model.trainable = False
    tf_image_processor = EfficientNetImageProcessor(
        do_center_crop=True,
        do_rescale=False,
        do_normalize=False,
        include_top=False,
        resample=Image.BILINEAR,
    )
    image = tf_image_processor(images=prepare_img(), return_tensors="tf", data_format="channels_last")["pixel_values"]
    text = tok(tf.constant(["A picture of a cat"]))

    image_features = original_model.image_encoder(image, training=False)
    text_features = original_model.text_encoder(text, training=False)

    image_features = tf.nn.l2_normalize(image_features, axis=-1)
    text_features = tf.nn.l2_normalize(text_features, axis=-1)

    
    if not np.allclose(image_features, hf_image_features, atol=1e-3):
        raise ValueError("The predicted image features are not the same.")
    if not np.allclose(text_features, hf_text_features, atol=1e-3):
        raise ValueError("The predicted text features are not the same.")
    print("Model outputs match!")

    if save_model:
        
        if not os.path.isdir(pytorch_dump_folder_path):
            os.mkdir(pytorch_dump_folder_path)
        
        hf_model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_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__":
    parser = argparse.ArgumentParser()
    
    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")

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