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 dpt 3 1 checkpoints from the midas repository url https github comislorgmidas import argparse from pathlib import path import requests import torch from pil import image from transformers import beitconfig dptconfig dptfordepthestimation dptimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getdptconfigmodelname hiddensize 768 numhiddenlayers 12 numattentionheads 12 intermediatesize 3072 outfeatures stage3 stage6 stage9 stage12 beitbase384 uses 2 5 8 11 if large in modelname hiddensize 1024 numhiddenlayers 24 numattentionheads 16 intermediatesize 4096 outfeatures stage6 stage12 stage18 stage24 beitlarge512 uses 5 11 17 23 if 512 in modelname imagesize 512 elif 384 in modelname imagesize 384 else raise valueerrormodel not supported backboneconfig beitconfig imagesizeimagesize numhiddenlayersnumhiddenlayers hiddensizehiddensize intermediatesizeintermediatesize numattentionheadsnumattentionheads userelativepositionbiastrue reshapehiddenstatesfalse outfeaturesoutfeatures neckhiddensizes 256 512 1024 1024 if large in modelname else 96 192 384 768 config dptconfigbackboneconfigbackboneconfig neckhiddensizesneckhiddensizes return config imagesize here we list all keys to be renamed original name on the left our name on the right def createrenamekeysconfig renamekeys fmt off stem renamekeys appendpretrained model clstoken backbone embeddings clstoken renamekeys appendpretrained model patchembed proj weight backbone embeddings patchembeddings projection weight renamekeys appendpretrained model patchembed proj bias backbone embeddings patchembeddings projection bias transfomer encoder for i in rangeconfig backboneconfig numhiddenlayers renamekeys appendfpretrained model blocks i gamma1 fbackbone encoder layer i lambda1 renamekeys appendfpretrained model blocks i gamma2 fbackbone encoder layer i lambda2 renamekeys appendfpretrained model blocks i norm1 weight fbackbone encoder layer i layernormbefore weight renamekeys appendfpretrained model blocks i norm1 bias fbackbone encoder layer i layernormbefore bias renamekeys appendfpretrained model blocks i norm2 weight fbackbone encoder layer i layernormafter weight renamekeys appendfpretrained model blocks i norm2 bias fbackbone encoder layer i layernormafter bias renamekeys appendfpretrained model blocks i mlp fc1 weight fbackbone encoder layer i intermediate dense weight renamekeys appendfpretrained model blocks i mlp fc1 bias fbackbone encoder layer i intermediate dense bias renamekeys appendfpretrained model blocks i mlp fc2 weight fbackbone encoder layer i output dense weight renamekeys appendfpretrained model blocks i mlp fc2 bias fbackbone encoder layer i output dense bias renamekeys appendfpretrained model blocks i attn proj weight fbackbone encoder layer i attention output dense weight renamekeys appendfpretrained model blocks i attn proj bias fbackbone encoder layer i attention output dense bias renamekeys appendfpretrained model blocks i attn relativepositionbiastable fbackbone encoder layer i attention attention relativepositionbias relativepositionbiastable renamekeys appendfpretrained model blocks i attn relativepositionindex fbackbone encoder layer i attention attention relativepositionbias relativepositionindex activation postprocessing readout projections resize blocks for i in range4 renamekeys appendfpretrained actpostprocessi1 0 project 0 weight fneck reassemblestage readoutprojects i 0 weight renamekeys appendfpretrained actpostprocessi1 0 project 0 bias fneck reassemblestage readoutprojects i 0 bias renamekeys appendfpretrained actpostprocessi1 3 weight fneck reassemblestage layers i projection weight renamekeys appendfpretrained actpostprocessi1 3 bias fneck reassemblestage layers i projection bias if i 2 renamekeys appendfpretrained actpostprocessi1 4 weight fneck reassemblestage layers i resize weight renamekeys appendfpretrained actpostprocessi1 4 bias fneck reassemblestage layers i resize bias refinenet tricky here mapping 1 3 2 2 3 1 4 0 for i in range1 5 j mappingi renamekeys appendfscratch refineneti outconv weight fneck fusionstage layers j projection weight renamekeys appendfscratch refineneti outconv bias fneck fusionstage layers j projection bias renamekeys appendfscratch refineneti resconfunit1 conv1 weight fneck fusionstage layers j residuallayer1 convolution1 weight renamekeys appendfscratch refineneti resconfunit1 conv1 bias fneck fusionstage layers j residuallayer1 convolution1 bias renamekeys appendfscratch refineneti resconfunit1 conv2 weight fneck fusionstage layers j residuallayer1 convolution2 weight renamekeys appendfscratch refineneti resconfunit1 conv2 bias fneck fusionstage layers j residuallayer1 convolution2 bias renamekeys appendfscratch refineneti resconfunit2 conv1 weight fneck fusionstage layers j residuallayer2 convolution1 weight renamekeys appendfscratch refineneti resconfunit2 conv1 bias fneck fusionstage layers j residuallayer2 convolution1 bias renamekeys appendfscratch refineneti resconfunit2 conv2 weight fneck fusionstage layers j residuallayer2 convolution2 weight renamekeys appendfscratch refineneti resconfunit2 conv2 bias fneck fusionstage layers j residuallayer2 convolution2 bias scratch convolutions for i in range4 renamekeys appendfscratch layeri1rn weight fneck convs i weight head for i in range0 5 2 renamekeys appendfscratch outputconv i weight fhead head i weight renamekeys appendfscratch outputconv i bias fhead head i bias return renamekeys def removeignorekeysstatedict ignorekeys pretrained model head weight pretrained model head bias for k in ignorekeys statedict popk none we split up the matrix of each encoder layer into queries keys and values def readinqkvstatedict config hiddensize config backboneconfig hiddensize for i in rangeconfig backboneconfig numhiddenlayers read in weights bias of input projection layer in original implementation this is a single matrix bias inprojweight statedict popfpretrained model blocks i attn qkv weight qbias statedict popfpretrained model blocks i attn qbias vbias statedict popfpretrained model blocks i attn vbias next add query keys and values in that order to the state dict statedictfbackbone encoder layer i attention attention query weight inprojweight hiddensize statedictfbackbone encoder layer i attention attention query bias qbias statedictfbackbone encoder layer i attention attention key weight inprojweight hiddensize hiddensize 2 statedictfbackbone encoder layer i attention attention value weight inprojweighthiddensize statedictfbackbone encoder layer i attention attention value bias vbias def renamekeydct old new val dct popold dctnew val 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 torch nograd def convertdptcheckpointmodelname pytorchdumpfolderpath pushtohub nametourl dptbeitlarge512 https github comislorgmidasreleasesdownloadv31dptbeitlarge512 pt dptbeitlarge384 https github comislorgmidasreleasesdownloadv31dptbeitlarge384 pt dptbeitbase384 https github comislorgmidasreleasesdownloadv31dptbeitbase384 pt define dpt configuration based on url checkpointurl nametourlmodelname config imagesize getdptconfigmodelname load original statedict from url statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu remove certain keys removeignorekeysstatedict rename keys renamekeys createrenamekeysconfig for src dest in renamekeys renamekeystatedict src dest read in qkv matrices readinqkvstatedict config load huggingface model model dptfordepthestimationconfig missingkeys unexpectedkeys model loadstatedictstatedict strictfalse printmissing keys missingkeys printunexpected keys unexpectedkeys assert missingkeys assert unexpectedkeys pretrained model fcnorm weight pretrained model fcnorm bias model eval check outputs on an image processor dptimageprocessor sizeheight imagesize width imagesize keepaspectratiotrue ensuremultipleof32 image prepareimg pixelvalues processorimage returntensorspt pixelvalues printfirst values of pixel values pixelvalues0 0 3 3 printmean of pixel values pixelvalues mean item printshape of pixel values pixelvalues shape import requests from pil import image from torchvision import transforms url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw transforms transforms compose transforms resizeimagesize imagesize transforms totensor pixelvalues transformsimage unsqueeze0 forward pass with torch nograd outputs modelpixelvalues predicteddepth outputs predicteddepth printshape of predicted depth predicteddepth shape printfirst values of predicted depth predicteddepth0 3 3 assert logits todo there s still a small difference with the original logits if modelname dptbeitlarge512 ok checked expectedshape torch size1 512 512 expectedslice torch tensor 2804 6260 2792 5708 2812 9263 2772 0288 2780 1118 2796 2529 2748 1094 2766 6558 2766 9834 elif modelname dptbeitlarge384 ok checked expectedshape torch size1 384 384 expectedslice torch tensor 1783 2273 1780 5729 1792 6453 1759 9817 1765 5359 1778 5002 1739 1633 1754 7903 1757 1990 elif modelname dptbeitbase384 ok checked expectedshape torch size1 384 384 expectedslice torch tensor 2898 4482 2891 3750 2904 8079 2858 6685 2877 2615 2894 4507 2842 1235 2854 1023 2861 6328 assert predicteddepth shape torch sizeexpectedshape assert torch allclosepredicteddepth0 3 3 expectedslice printlooks ok if pytorchdumpfolderpath is not none pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model and processor to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub printpushing model and processor to hub model pushtohubrepoidfnielsrmodelname processor pushtohubrepoidfnielsrmodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultdptbeitlarge512 typestr choicesdptbeitlarge512 dptbeitlarge384 dptbeitbase384 helpname of the model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the output pytorch model directory parser addargument pushtohub actionstoretrue helpwhether to push the model to the hub after conversion args parser parseargs convertdptcheckpointargs modelname args pytorchdumpfolderpath 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 dpt 3 1 checkpoints from the midas repository url https github com isl org midas beit base 384 uses 2 5 8 11 beit large 512 uses 5 11 17 23 here we list all keys to be renamed original name on the left our name on the right fmt off stem transfomer encoder activation postprocessing readout projections resize blocks refinenet tricky here scratch convolutions head we split up the matrix of each encoder layer into queries keys and values read in weights bias of input projection layer in original implementation this is a single matrix bias next add query keys and values in that order to the state dict we will verify our results on an image of cute cats copy paste tweak model s weights to our dpt structure define dpt configuration based on url load original state_dict from url remove certain keys rename keys read in qkv matrices load huggingface model assert unexpected_keys pretrained model fc_norm weight pretrained model fc_norm bias check outputs on an image forward pass assert logits todo there s still a small difference with the original logits ok checked ok checked ok checked required parameters
import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import BeitConfig, DPTConfig, DPTForDepthEstimation, DPTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)


def get_dpt_config(model_name):
    hidden_size = 768
    num_hidden_layers = 12
    num_attention_heads = 12
    intermediate_size = 3072
    out_features = ["stage3", "stage6", "stage9", "stage12"]  

    if "large" in model_name:
        hidden_size = 1024
        num_hidden_layers = 24
        num_attention_heads = 16
        intermediate_size = 4096
        out_features = ["stage6", "stage12", "stage18", "stage24"]  

    if "512" in model_name:
        image_size = 512
    elif "384" in model_name:
        image_size = 384
    else:
        raise ValueError("Model not supported")

    backbone_config = BeitConfig(
        image_size=image_size,
        num_hidden_layers=num_hidden_layers,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_attention_heads=num_attention_heads,
        use_relative_position_bias=True,
        reshape_hidden_states=False,
        out_features=out_features,
    )

    neck_hidden_sizes = [256, 512, 1024, 1024] if "large" in model_name else [96, 192, 384, 768]
    config = DPTConfig(backbone_config=backbone_config, neck_hidden_sizes=neck_hidden_sizes)

    return config, image_size



def create_rename_keys(config):
    rename_keys = []

    
    
    rename_keys.append(("pretrained.model.cls_token", "backbone.embeddings.cls_token"))
    rename_keys.append(("pretrained.model.patch_embed.proj.weight", "backbone.embeddings.patch_embeddings.projection.weight"))
    rename_keys.append(("pretrained.model.patch_embed.proj.bias", "backbone.embeddings.patch_embeddings.projection.bias"))

    
    for i in range(config.backbone_config.num_hidden_layers):
        rename_keys.append((f"pretrained.model.blocks.{i}.gamma_1", f"backbone.encoder.layer.{i}.lambda_1"))
        rename_keys.append((f"pretrained.model.blocks.{i}.gamma_2", f"backbone.encoder.layer.{i}.lambda_2"))
        rename_keys.append((f"pretrained.model.blocks.{i}.norm1.weight", f"backbone.encoder.layer.{i}.layernorm_before.weight"))
        rename_keys.append((f"pretrained.model.blocks.{i}.norm1.bias", f"backbone.encoder.layer.{i}.layernorm_before.bias"))
        rename_keys.append((f"pretrained.model.blocks.{i}.norm2.weight", f"backbone.encoder.layer.{i}.layernorm_after.weight"))
        rename_keys.append((f"pretrained.model.blocks.{i}.norm2.bias", f"backbone.encoder.layer.{i}.layernorm_after.bias"))
        rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.weight", f"backbone.encoder.layer.{i}.intermediate.dense.weight"))
        rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc1.bias", f"backbone.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.weight", f"backbone.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"pretrained.model.blocks.{i}.mlp.fc2.bias", f"backbone.encoder.layer.{i}.output.dense.bias"))
        rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.weight", f"backbone.encoder.layer.{i}.attention.output.dense.weight"))
        rename_keys.append((f"pretrained.model.blocks.{i}.attn.proj.bias", f"backbone.encoder.layer.{i}.attention.output.dense.bias"))
        rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_bias_table", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_bias_table"))
        rename_keys.append((f"pretrained.model.blocks.{i}.attn.relative_position_index", f"backbone.encoder.layer.{i}.attention.attention.relative_position_bias.relative_position_index"))

    
    for i in range(4):
        rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.weight", f"neck.reassemble_stage.readout_projects.{i}.0.weight"))
        rename_keys.append((f"pretrained.act_postprocess{i+1}.0.project.0.bias", f"neck.reassemble_stage.readout_projects.{i}.0.bias"))

        rename_keys.append((f"pretrained.act_postprocess{i+1}.3.weight", f"neck.reassemble_stage.layers.{i}.projection.weight"))
        rename_keys.append((f"pretrained.act_postprocess{i+1}.3.bias", f"neck.reassemble_stage.layers.{i}.projection.bias"))

        if i != 2:
            rename_keys.append((f"pretrained.act_postprocess{i+1}.4.weight", f"neck.reassemble_stage.layers.{i}.resize.weight"))
            rename_keys.append((f"pretrained.act_postprocess{i+1}.4.bias", f"neck.reassemble_stage.layers.{i}.resize.bias"))

    
    mapping = {1:3, 2:2, 3:1, 4:0}

    for i in range(1, 5):
        j = mapping[i]
        rename_keys.append((f"scratch.refinenet{i}.out_conv.weight", f"neck.fusion_stage.layers.{j}.projection.weight"))
        rename_keys.append((f"scratch.refinenet{i}.out_conv.bias", f"neck.fusion_stage.layers.{j}.projection.bias"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.weight"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution1.bias"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.weight"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit1.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer1.convolution2.bias"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.weight"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv1.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution1.bias"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.weight", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.weight"))
        rename_keys.append((f"scratch.refinenet{i}.resConfUnit2.conv2.bias", f"neck.fusion_stage.layers.{j}.residual_layer2.convolution2.bias"))

    
    for i in range(4):
        rename_keys.append((f"scratch.layer{i+1}_rn.weight", f"neck.convs.{i}.weight"))

    
    for i in range(0, 5, 2):
        rename_keys.append((f"scratch.output_conv.{i}.weight", f"head.head.{i}.weight"))
        rename_keys.append((f"scratch.output_conv.{i}.bias", f"head.head.{i}.bias"))

    return rename_keys


def remove_ignore_keys_(state_dict):
    ignore_keys = ["pretrained.model.head.weight", "pretrained.model.head.bias"]
    for k in ignore_keys:
        state_dict.pop(k, None)



def read_in_q_k_v(state_dict, config):
    hidden_size = config.backbone_config.hidden_size
    for i in range(config.backbone_config.num_hidden_layers):
        
        in_proj_weight = state_dict.pop(f"pretrained.model.blocks.{i}.attn.qkv.weight")
        q_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.q_bias")
        v_bias = state_dict.pop(f"pretrained.model.blocks.{i}.attn.v_bias")
        
        state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[:hidden_size, :]
        state_dict[f"backbone.encoder.layer.{i}.attention.attention.query.bias"] = q_bias
        state_dict[f"backbone.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
            hidden_size : hidden_size * 2, :
        ]
        state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[-hidden_size:, :]
        state_dict[f"backbone.encoder.layer.{i}.attention.attention.value.bias"] = v_bias


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val



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


@torch.no_grad()
def convert_dpt_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub):
    

    name_to_url = {
        "dpt-beit-large-512": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt",
        "dpt-beit-large-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_384.pt",
        "dpt-beit-base-384": "https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_base_384.pt",
    }

    
    checkpoint_url = name_to_url[model_name]
    config, image_size = get_dpt_config(model_name)
    
    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
    
    remove_ignore_keys_(state_dict)
    
    rename_keys = create_rename_keys(config)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    
    read_in_q_k_v(state_dict, config)

    
    model = DPTForDepthEstimation(config)
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    print("Missing keys:", missing_keys)
    print("Unexpected keys:", unexpected_keys)
    assert missing_keys == []
    
    model.eval()

    
    processor = DPTImageProcessor(
        size={"height": image_size, "width": image_size}, keep_aspect_ratio=True, ensure_multiple_of=32
    )

    image = prepare_img()
    pixel_values = processor(image, return_tensors="pt").pixel_values

    print("First values of pixel values:", pixel_values[0, 0, :3, :3])
    print("Mean of pixel values:", pixel_values.mean().item())
    print("Shape of pixel values:", pixel_values.shape)

    import requests
    from PIL import Image
    from torchvision import transforms

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

    transforms = transforms.Compose(
        [
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ]
    )
    pixel_values = transforms(image).unsqueeze(0)

    
    with torch.no_grad():
        outputs = model(pixel_values)

    predicted_depth = outputs.predicted_depth

    print("Shape of predicted depth:", predicted_depth.shape)
    print("First values of predicted depth:", predicted_depth[0, :3, :3])

    
    
    if model_name == "dpt-beit-large-512":
        
        expected_shape = torch.Size([1, 512, 512])
        expected_slice = torch.tensor(
            [[2804.6260, 2792.5708, 2812.9263], [2772.0288, 2780.1118, 2796.2529], [2748.1094, 2766.6558, 2766.9834]]
        )
    elif model_name == "dpt-beit-large-384":
        
        expected_shape = torch.Size([1, 384, 384])
        expected_slice = torch.tensor(
            [[1783.2273, 1780.5729, 1792.6453], [1759.9817, 1765.5359, 1778.5002], [1739.1633, 1754.7903, 1757.1990]],
        )
    elif model_name == "dpt-beit-base-384":
        
        expected_shape = torch.Size([1, 384, 384])
        expected_slice = torch.tensor(
            [[2898.4482, 2891.3750, 2904.8079], [2858.6685, 2877.2615, 2894.4507], [2842.1235, 2854.1023, 2861.6328]],
        )

    assert predicted_depth.shape == torch.Size(expected_shape)
    assert torch.allclose(predicted_depth[0, :3, :3], expected_slice)
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
        print(f"Saving model and processor to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing model and processor to hub...")
        model.push_to_hub(repo_id=f"nielsr/{model_name}")
        processor.push_to_hub(repo_id=f"nielsr/{model_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="dpt-beit-large-512",
        type=str,
        choices=["dpt-beit-large-512", "dpt-beit-large-384", "dpt-beit-base-384"],
        help="Name of the model you'd like to convert.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument(
        "--push_to_hub",
        action="store_true",
        help="Whether to push the model to the hub after conversion.",
    )

    args = parser.parse_args()
    convert_dpt_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)