codingutf8 2022 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 glpn checkpoints import argparse from collections import ordereddict from pathlib import path import requests import torch from pil import image from transformers import glpnconfig glpnfordepthestimation glpnimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def renamekeysstatedict newstatedict ordereddict for key value in statedict items if key startswithmodule encoder key key replacemodule encoder glpn encoder if key startswithmodule decoder key key replacemodule decoder decoder stages if patchembed in key replace for example patchembed1 by patchembeddings 0 idx keykey findpatchembed lenpatchembed key key replacefpatchembedidx fpatchembeddings intidx1 if norm in key key key replacenorm layernorm if glpn encoder layernorm in key replace for example layernorm1 by layernorm 0 idx keykey findglpn encoder layernorm lenglpn encoder layernorm key key replaceflayernormidx flayernorm intidx1 if layernorm1 in key key key replacelayernorm1 layernorm1 if layernorm2 in key key key replacelayernorm2 layernorm2 if block in key replace for example block1 by block 0 idx keykey findblock lenblock key key replacefblockidx fblock intidx1 if attn q in key key key replaceattn q attention self query if attn proj in key key key replaceattn proj attention output dense if attn in key key key replaceattn attention self if fc1 in key key key replacefc1 dense1 if fc2 in key key key replacefc2 dense2 if linearpred in key key key replacelinearpred classifier if linearfuse in key key key replacelinearfuse conv linearfuse key key replacelinearfuse bn batchnorm if linearc in key replace for example linearc4 by linearc 3 idx keykey findlinearc lenlinearc key key replaceflinearcidx flinearc intidx1 if botconv in key key key replacebotconv 0 convolution if skipconv1 in key key key replaceskipconv1 1 convolution if skipconv2 in key key key replaceskipconv2 2 convolution if fusion1 in key key key replacefusion1 1 fusion if fusion2 in key key key replacefusion2 2 fusion if fusion3 in key key key replacefusion3 3 fusion if fusion in key and conv in key key key replaceconv convolutionallayer if key startswithmodule lastlayerdepth key key replacemodule lastlayerdepth head head newstatedictkey value return newstatedict def readinkvstatedict config for each of the encoder blocks for i in rangeconfig numencoderblocks for j in rangeconfig depthsi read in weights bias of keys and values which is a single matrix in the original implementation kvweight statedict popfglpn encoder block i j attention self kv weight kvbias statedict popfglpn encoder block i j attention self kv bias next add keys and values in that order to the state dict statedictfglpn encoder block i j attention self key weight kvweight config hiddensizesi statedictfglpn encoder block i j attention self key bias kvbias config hiddensizesi statedictfglpn encoder block i j attention self value weight kvweight config hiddensizesi statedictfglpn encoder block i j attention self value bias kvbiasconfig hiddensizesi we will verify our results on a coco image def prepareimg url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw return image torch nograd def convertglpncheckpointcheckpointpath pytorchdumpfolderpath pushtohubfalse modelnamenone load glpn configuration segformerb4 size config glpnconfighiddensizes64 128 320 512 decoderhiddensize64 depths3 8 27 3 load image processor only resize rescale imageprocessor glpnimageprocessor prepare image image prepareimg pixelvalues imageprocessorimagesimage returntensorspt pixelvalues logger infoconverting model load original state dict statedict torch loadcheckpointpath maplocationtorch devicecpu rename keys statedict renamekeysstatedict key and value matrices need special treatment readinkvstatedict config create huggingface model and load state dict model glpnfordepthestimationconfig model loadstatedictstatedict model eval forward pass outputs modelpixelvalues predicteddepth outputs predicteddepth verify output if modelname is not none if nyu in modelname expectedslice torch tensor 4 4147 4 0873 4 0673 3 7890 3 2881 3 1525 3 7674 3 5423 3 4913 elif kitti in modelname expectedslice torch tensor 3 4291 2 7865 2 5151 3 2841 2 7021 2 3502 3 1147 2 4625 2 2481 else raise valueerrorfunknown model name modelname expectedshape torch size1 480 640 assert predicteddepth shape expectedshape assert torch allclosepredicteddepth0 3 3 expectedslice atol1e4 printlooks ok finally push to hub if required if pushtohub logger infopushing model and image processor to the hub model pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd model usetempdirtrue imageprocessor pushtohub repopathornamepathpytorchdumpfolderpath modelname organizationnielsr commitmessageadd image processor usetempdirtrue if name main parser argparse argumentparser parser addargument checkpointpath defaultnone typestr helppath to the original pytorch checkpoint pth file parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model parser addargument pushtohub actionstoretrue helpwhether to upload the model to the huggingface hub parser addargument modelname defaultglpnkitti typestr helpname of the model in case you re pushing to the hub args parser parseargs convertglpncheckpointargs checkpointpath args pytorchdumpfolderpath args pushtohub args modelname coding utf 8 2022 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 glpn checkpoints replace for example patch_embed1 by patch_embeddings 0 replace for example layer_norm1 by layer_norm 0 replace for example block1 by block 0 replace for example linear_c4 by linear_c 3 for each of the encoder blocks read in weights bias of keys and values which is a single matrix in the original implementation next add keys and values in that order to the state dict we will verify our results on a coco image copy paste tweak model s weights to our glpn structure load glpn configuration segformer b4 size load image processor only resize rescale prepare image load original state dict rename keys key and value matrices need special treatment create huggingface model and load state dict forward pass verify output finally push to hub if required
import argparse
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import GLPNConfig, GLPNForDepthEstimation, GLPNImageProcessor
from transformers.utils import logging


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


def rename_keys(state_dict):
    new_state_dict = OrderedDict()
    for key, value in state_dict.items():
        if key.startswith("module.encoder"):
            key = key.replace("module.encoder", "glpn.encoder")
        if key.startswith("module.decoder"):
            key = key.replace("module.decoder", "decoder.stages")
        if "patch_embed" in key:
            
            idx = key[key.find("patch_embed") + len("patch_embed")]
            key = key.replace(f"patch_embed{idx}", f"patch_embeddings.{int(idx)-1}")
        if "norm" in key:
            key = key.replace("norm", "layer_norm")
        if "glpn.encoder.layer_norm" in key:
            
            idx = key[key.find("glpn.encoder.layer_norm") + len("glpn.encoder.layer_norm")]
            key = key.replace(f"layer_norm{idx}", f"layer_norm.{int(idx)-1}")
        if "layer_norm1" in key:
            key = key.replace("layer_norm1", "layer_norm_1")
        if "layer_norm2" in key:
            key = key.replace("layer_norm2", "layer_norm_2")
        if "block" in key:
            
            idx = key[key.find("block") + len("block")]
            key = key.replace(f"block{idx}", f"block.{int(idx)-1}")
        if "attn.q" in key:
            key = key.replace("attn.q", "attention.self.query")
        if "attn.proj" in key:
            key = key.replace("attn.proj", "attention.output.dense")
        if "attn" in key:
            key = key.replace("attn", "attention.self")
        if "fc1" in key:
            key = key.replace("fc1", "dense1")
        if "fc2" in key:
            key = key.replace("fc2", "dense2")
        if "linear_pred" in key:
            key = key.replace("linear_pred", "classifier")
        if "linear_fuse" in key:
            key = key.replace("linear_fuse.conv", "linear_fuse")
            key = key.replace("linear_fuse.bn", "batch_norm")
        if "linear_c" in key:
            
            idx = key[key.find("linear_c") + len("linear_c")]
            key = key.replace(f"linear_c{idx}", f"linear_c.{int(idx)-1}")
        if "bot_conv" in key:
            key = key.replace("bot_conv", "0.convolution")
        if "skip_conv1" in key:
            key = key.replace("skip_conv1", "1.convolution")
        if "skip_conv2" in key:
            key = key.replace("skip_conv2", "2.convolution")
        if "fusion1" in key:
            key = key.replace("fusion1", "1.fusion")
        if "fusion2" in key:
            key = key.replace("fusion2", "2.fusion")
        if "fusion3" in key:
            key = key.replace("fusion3", "3.fusion")
        if "fusion" in key and "conv" in key:
            key = key.replace("conv", "convolutional_layer")
        if key.startswith("module.last_layer_depth"):
            key = key.replace("module.last_layer_depth", "head.head")
        new_state_dict[key] = value

    return new_state_dict


def read_in_k_v(state_dict, config):
    
    for i in range(config.num_encoder_blocks):
        for j in range(config.depths[i]):
            
            kv_weight = state_dict.pop(f"glpn.encoder.block.{i}.{j}.attention.self.kv.weight")
            kv_bias = state_dict.pop(f"glpn.encoder.block.{i}.{j}.attention.self.kv.bias")
            
            state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.key.weight"] = kv_weight[
                : config.hidden_sizes[i], :
            ]
            state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.key.bias"] = kv_bias[: config.hidden_sizes[i]]
            state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.value.weight"] = kv_weight[
                config.hidden_sizes[i] :, :
            ]
            state_dict[f"glpn.encoder.block.{i}.{j}.attention.self.value.bias"] = kv_bias[config.hidden_sizes[i] :]



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

    return image


@torch.no_grad()
def convert_glpn_checkpoint(checkpoint_path, pytorch_dump_folder_path, push_to_hub=False, model_name=None):
    

    
    config = GLPNConfig(hidden_sizes=[64, 128, 320, 512], decoder_hidden_size=64, depths=[3, 8, 27, 3])

    
    image_processor = GLPNImageProcessor()

    
    image = prepare_img()
    pixel_values = image_processor(images=image, return_tensors="pt").pixel_values

    logger.info("Converting model...")

    
    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))

    
    state_dict = rename_keys(state_dict)

    
    read_in_k_v(state_dict, config)

    
    model = GLPNForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    
    outputs = model(pixel_values)
    predicted_depth = outputs.predicted_depth

    
    if model_name is not None:
        if "nyu" in model_name:
            expected_slice = torch.tensor(
                [[4.4147, 4.0873, 4.0673], [3.7890, 3.2881, 3.1525], [3.7674, 3.5423, 3.4913]]
            )
        elif "kitti" in model_name:
            expected_slice = torch.tensor(
                [[3.4291, 2.7865, 2.5151], [3.2841, 2.7021, 2.3502], [3.1147, 2.4625, 2.2481]]
            )
        else:
            raise ValueError(f"Unknown model name: {model_name}")

        expected_shape = torch.Size([1, 480, 640])

        assert predicted_depth.shape == expected_shape
        assert torch.allclose(predicted_depth[0, :3, :3], expected_slice, atol=1e-4)
        print("Looks ok!")

    
    if push_to_hub:
        logger.info("Pushing model and image processor to the hub...")
        model.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add model",
            use_temp_dir=True,
        )
        image_processor.push_to_hub(
            repo_path_or_name=Path(pytorch_dump_folder_path, model_name),
            organization="nielsr",
            commit_message="Add image processor",
            use_temp_dir=True,
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        help="Path to the original PyTorch checkpoint (.pth file).",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", help="Whether to upload the model to the HuggingFace hub."
    )
    parser.add_argument(
        "--model_name",
        default="glpn-kitti",
        type=str,
        help="Name of the model in case you're pushing to the hub.",
    )
    args = parser.parse_args()
    convert_glpn_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)