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 detr checkpoints with native transformers backbone import argparse import json from pathlib import path import requests import torch from huggingfacehub import hfhubdownload from pil import image from transformers import detrconfig detrforobjectdetection detrforsegmentation detrimageprocessor resnetconfig from transformers utils import logging logging setverbosityinfo logger logging getloggername def getdetrconfigmodelname initialize config if resnet50 in modelname backboneconfig resnetconfig frompretrainedmicrosoftresnet50 elif resnet101 in modelname backboneconfig resnetconfig frompretrainedmicrosoftresnet101 else raise valueerrormodel name should include either resnet50 or resnet101 config detrconfigusetimmbackbonefalse backboneconfigbackboneconfig set label attributes ispanoptic panoptic in modelname if ispanoptic config numlabels 250 else config numlabels 91 repoid huggingfacelabelfiles filename cocodetectionid2label json id2label json loadopenhfhubdownloadrepoid filename repotypedataset r id2label intk v for k v in id2label items config id2label id2label config label2id v k for k v in id2label items return config ispanoptic def createrenamekeysconfig here we list all keys to be renamed original name on the left our name on the right renamekeys stem fmt off renamekeys appendbackbone 0 body conv1 weight backbone convencoder model embedder embedder convolution weight renamekeys appendbackbone 0 body bn1 weight backbone convencoder model embedder embedder normalization weight renamekeys appendbackbone 0 body bn1 bias backbone convencoder model embedder embedder normalization bias renamekeys appendbackbone 0 body bn1 runningmean backbone convencoder model embedder embedder normalization runningmean renamekeys appendbackbone 0 body bn1 runningvar backbone convencoder model embedder embedder normalization runningvar stages for stageidx in rangelenconfig backboneconfig depths for layeridx in rangeconfig backboneconfig depthsstageidx shortcut if layeridx 0 renamekeys append fbackbone 0 body layerstageidx 1 layeridx downsample 0 weight fbackbone convencoder model encoder stages stageidx layers layeridx shortcut convolution weight renamekeys append fbackbone 0 body layerstageidx 1 layeridx downsample 1 weight fbackbone convencoder model encoder stages stageidx layers layeridx shortcut normalization weight renamekeys append fbackbone 0 body layerstageidx 1 layeridx downsample 1 bias fbackbone convencoder model encoder stages stageidx layers layeridx shortcut normalization bias renamekeys append fbackbone 0 body layerstageidx 1 layeridx downsample 1 runningmean fbackbone convencoder model encoder stages stageidx layers layeridx shortcut normalization runningmean renamekeys append fbackbone 0 body layerstageidx 1 layeridx downsample 1 runningvar fbackbone convencoder model encoder stages stageidx layers layeridx shortcut normalization runningvar 3 convs for i in range3 renamekeys append fbackbone 0 body layerstageidx 1 layeridx convi1 weight fbackbone convencoder model encoder stages stageidx layers layeridx layer i convolution weight renamekeys append fbackbone 0 body layerstageidx 1 layeridx bni1 weight fbackbone convencoder model encoder stages stageidx layers layeridx layer i normalization weight renamekeys append fbackbone 0 body layerstageidx 1 layeridx bni1 bias fbackbone convencoder model encoder stages stageidx layers layeridx layer i normalization bias renamekeys append fbackbone 0 body layerstageidx 1 layeridx bni1 runningmean fbackbone convencoder model encoder stages stageidx layers layeridx layer i normalization runningmean renamekeys append fbackbone 0 body layerstageidx 1 layeridx bni1 runningvar fbackbone convencoder model encoder stages stageidx layers layeridx layer i normalization runningvar fmt on for i in rangeconfig encoderlayers encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys append ftransformer encoder layers i selfattn outproj weight fencoder layers i selfattn outproj weight renamekeys append ftransformer encoder layers i selfattn outproj bias fencoder layers i selfattn outproj bias renamekeys appendftransformer encoder layers i linear1 weight fencoder layers i fc1 weight renamekeys appendftransformer encoder layers i linear1 bias fencoder layers i fc1 bias renamekeys appendftransformer encoder layers i linear2 weight fencoder layers i fc2 weight renamekeys appendftransformer encoder layers i linear2 bias fencoder layers i fc2 bias renamekeys append ftransformer encoder layers i norm1 weight fencoder layers i selfattnlayernorm weight renamekeys append ftransformer encoder layers i norm1 bias fencoder layers i selfattnlayernorm bias renamekeys append ftransformer encoder layers i norm2 weight fencoder layers i finallayernorm weight renamekeys appendftransformer encoder layers i norm2 bias fencoder layers i finallayernorm bias decoder layers 2 times output projection 2 feedforward neural networks and 3 layernorms renamekeys append ftransformer decoder layers i selfattn outproj weight fdecoder layers i selfattn outproj weight renamekeys append ftransformer decoder layers i selfattn outproj bias fdecoder layers i selfattn outproj bias renamekeys append ftransformer decoder layers i multiheadattn outproj weight fdecoder layers i encoderattn outproj weight renamekeys append ftransformer decoder layers i multiheadattn outproj bias fdecoder layers i encoderattn outproj bias renamekeys appendftransformer decoder layers i linear1 weight fdecoder layers i fc1 weight renamekeys appendftransformer decoder layers i linear1 bias fdecoder layers i fc1 bias renamekeys appendftransformer decoder layers i linear2 weight fdecoder layers i fc2 weight renamekeys appendftransformer decoder layers i linear2 bias fdecoder layers i fc2 bias renamekeys append ftransformer decoder layers i norm1 weight fdecoder layers i selfattnlayernorm weight renamekeys append ftransformer decoder layers i norm1 bias fdecoder layers i selfattnlayernorm bias renamekeys append ftransformer decoder layers i norm2 weight fdecoder layers i encoderattnlayernorm weight renamekeys append ftransformer decoder layers i norm2 bias fdecoder layers i encoderattnlayernorm bias renamekeys append ftransformer decoder layers i norm3 weight fdecoder layers i finallayernorm weight renamekeys appendftransformer decoder layers i norm3 bias fdecoder layers i finallayernorm bias convolutional projection query embeddings layernorm of decoder class and bounding box heads renamekeys extend inputproj weight inputprojection weight inputproj bias inputprojection bias queryembed weight querypositionembeddings weight transformer decoder norm weight decoder layernorm weight transformer decoder norm bias decoder layernorm bias classembed weight classlabelsclassifier weight classembed bias classlabelsclassifier bias bboxembed layers 0 weight bboxpredictor layers 0 weight bboxembed layers 0 bias bboxpredictor layers 0 bias bboxembed layers 1 weight bboxpredictor layers 1 weight bboxembed layers 1 bias bboxpredictor layers 1 bias bboxembed layers 2 weight bboxpredictor layers 2 weight bboxembed layers 2 bias bboxpredictor layers 2 bias return renamekeys def renamekeystatedict old new val statedict popold statedictnew val def readinqkvstatedict ispanopticfalse prefix if ispanoptic prefix detr first transformer encoder for i in range6 read in weights bias of input projection layer in pytorch s multiheadattention this is a single matrix bias inprojweight statedict popfprefixtransformer encoder layers i selfattn inprojweight inprojbias statedict popfprefixtransformer encoder layers i selfattn inprojbias next add query keys and values in that order to the state dict statedictfencoder layers i selfattn qproj weight inprojweight 256 statedictfencoder layers i selfattn qproj bias inprojbias 256 statedictfencoder layers i selfattn kproj weight inprojweight256 512 statedictfencoder layers i selfattn kproj bias inprojbias256 512 statedictfencoder layers i selfattn vproj weight inprojweight256 statedictfencoder layers i selfattn vproj bias inprojbias256 next transformer decoder which is a bit more complex because it also includes crossattention for i in range6 read in weights bias of input projection layer of selfattention inprojweight statedict popfprefixtransformer decoder layers i selfattn inprojweight inprojbias statedict popfprefixtransformer decoder layers i selfattn inprojbias next add query keys and values in that order to the state dict statedictfdecoder layers i selfattn qproj weight inprojweight 256 statedictfdecoder layers i selfattn qproj bias inprojbias 256 statedictfdecoder layers i selfattn kproj weight inprojweight256 512 statedictfdecoder layers i selfattn kproj bias inprojbias256 512 statedictfdecoder layers i selfattn vproj weight inprojweight256 statedictfdecoder layers i selfattn vproj bias inprojbias256 read in weights bias of input projection layer of crossattention inprojweightcrossattn statedict pop fprefixtransformer decoder layers i multiheadattn inprojweight inprojbiascrossattn statedict popfprefixtransformer decoder layers i multiheadattn inprojbias next add query keys and values in that order of crossattention to the state dict statedictfdecoder layers i encoderattn qproj weight inprojweightcrossattn 256 statedictfdecoder layers i encoderattn qproj bias inprojbiascrossattn 256 statedictfdecoder layers i encoderattn kproj weight inprojweightcrossattn256 512 statedictfdecoder layers i encoderattn kproj bias inprojbiascrossattn256 512 statedictfdecoder layers i encoderattn vproj weight inprojweightcrossattn256 statedictfdecoder layers i encoderattn vproj bias inprojbiascrossattn256 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 convertdetrcheckpointmodelname pytorchdumpfolderpathnone pushtohubfalse load default config config ispanoptic getdetrconfigmodelname load original model from torch hub modelnametooriginalname detrresnet50 detrresnet50 detrresnet101 detrresnet101 logger infofconverting model modelname detr torch hub loadfacebookresearchdetr modelnametooriginalnamemodelname pretrainedtrue eval statedict detr statedict rename keys for src dest in createrenamekeysconfig if ispanoptic src detr src renamekeystatedict src dest query key and value matrices need special treatment readinqkvstatedict ispanopticispanoptic important we need to prepend a prefix to each of the base model keys as the head models use different attributes for them prefix detr model if ispanoptic else model for key in statedict copy keys if ispanoptic if key startswithdetr and not key startswithclasslabelsclassifier and not key startswithbboxpredictor val statedict popkey statedictdetr model key4 val elif classlabelsclassifier in key or bboxpredictor in key val statedict popkey statedictdetr key val elif key startswithbboxattention or key startswithmaskhead continue else val statedict popkey statedictprefix key val else if not key startswithclasslabelsclassifier and not key startswithbboxpredictor val statedict popkey statedictprefix key val finally create huggingface model and load state dict model detrforsegmentationconfig if ispanoptic else detrforobjectdetectionconfig model loadstatedictstatedict model eval verify our conversion on an image format cocopanoptic if ispanoptic else cocodetection processor detrimageprocessorformatformat encoding processorimagesprepareimg returntensorspt pixelvalues encodingpixelvalues originaloutputs detrpixelvalues outputs modelpixelvalues assert torch allcloseoutputs logits originaloutputspredlogits atol1e3 assert torch allcloseoutputs predboxes originaloutputspredboxes atol1e3 if ispanoptic assert torch allcloseoutputs predmasks originaloutputspredmasks atol1e4 printlooks ok if pytorchdumpfolderpath is not none save model and image processor logger infofsaving pytorch model and image processor to pytorchdumpfolderpath pathpytorchdumpfolderpath mkdirexistoktrue model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub upload model and image processor to the hub logger infouploading pytorch model and image processor to the hub model pushtohubfnielsrmodelname processor pushtohubfnielsrmodelname if name main parser argparse argumentparser parser addargument modelname defaultdetrresnet50 typestr choicesdetrresnet50 detrresnet101 helpname of the detr model you d like to convert parser addargument pytorchdumpfolderpath defaultnone typestr helppath to the folder to output pytorch model parser addargumentpushtohub actionstoretrue helpwhether to push the model to the hub or not args parser parseargs convertdetrcheckpointargs 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 detr checkpoints with native transformers backbone initialize config set label attributes here we list all keys to be renamed original name on the left our name on the right stem fmt off stages shortcut 3 convs fmt on encoder layers output projection 2 feedforward neural networks and 2 layernorms decoder layers 2 times output projection 2 feedforward neural networks and 3 layernorms convolutional projection query embeddings layernorm of decoder class and bounding box heads first transformer encoder read in weights bias of input projection layer in pytorch s multiheadattention this is a single matrix bias next add query keys and values in that order to the state dict next transformer decoder which is a bit more complex because it also includes cross attention read in weights bias of input projection layer of self attention next add query keys and values in that order to the state dict read in weights bias of input projection layer of cross attention next add query keys and values in that order of cross attention to the state dict we will verify our results on an image of cute cats copy paste tweak model s weights to our detr structure load default config load original model from torch hub rename keys query key and value matrices need special treatment important we need to prepend a prefix to each of the base model keys as the head models use different attributes for them finally create huggingface model and load state dict verify our conversion on an image save model and image processor upload model and image processor to the hub
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


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


def get_detr_config(model_name):
    
    if "resnet-50" in model_name:
        backbone_config = ResNetConfig.from_pretrained("microsoft/resnet-50")
    elif "resnet-101" in model_name:
        backbone_config = ResNetConfig.from_pretrained("microsoft/resnet-101")
    else:
        raise ValueError("Model name should include either resnet50 or resnet101")

    config = DetrConfig(use_timm_backbone=False, backbone_config=backbone_config)

    
    is_panoptic = "panoptic" in model_name
    if is_panoptic:
        config.num_labels = 250
    else:
        config.num_labels = 91
        repo_id = "huggingface/label-files"
        filename = "coco-detection-id2label.json"
        id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
        id2label = {int(k): v for k, v in id2label.items()}
        config.id2label = id2label
        config.label2id = {v: k for k, v in id2label.items()}

    return config, is_panoptic


def create_rename_keys(config):
    
    rename_keys = []

    
    
    rename_keys.append(("backbone.0.body.conv1.weight", "backbone.conv_encoder.model.embedder.embedder.convolution.weight"))
    rename_keys.append(("backbone.0.body.bn1.weight", "backbone.conv_encoder.model.embedder.embedder.normalization.weight"))
    rename_keys.append(("backbone.0.body.bn1.bias", "backbone.conv_encoder.model.embedder.embedder.normalization.bias"))
    rename_keys.append(("backbone.0.body.bn1.running_mean", "backbone.conv_encoder.model.embedder.embedder.normalization.running_mean"))
    rename_keys.append(("backbone.0.body.bn1.running_var", "backbone.conv_encoder.model.embedder.embedder.normalization.running_var"))
    
    for stage_idx in range(len(config.backbone_config.depths)):
        for layer_idx in range(config.backbone_config.depths[stage_idx]):
            
            if layer_idx == 0:
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var",
                    )
                )
            
            for i in range(3):
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean",
                    )
                )
                rename_keys.append(
                    (
                        f"backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var",
                        f"backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var",
                    )
                )
    

    for i in range(config.encoder_layers):
        
        rename_keys.append(
            (
                f"transformer.encoder.layers.{i}.self_attn.out_proj.weight",
                f"encoder.layers.{i}.self_attn.out_proj.weight",
            )
        )
        rename_keys.append(
            (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias")
        )
        rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight"))
        rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias"))
        rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight"))
        rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias"))
        rename_keys.append(
            (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight")
        )
        rename_keys.append(
            (f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias")
        )
        rename_keys.append(
            (f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight")
        )
        rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias"))
        
        rename_keys.append(
            (
                f"transformer.decoder.layers.{i}.self_attn.out_proj.weight",
                f"decoder.layers.{i}.self_attn.out_proj.weight",
            )
        )
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias")
        )
        rename_keys.append(
            (
                f"transformer.decoder.layers.{i}.multihead_attn.out_proj.weight",
                f"decoder.layers.{i}.encoder_attn.out_proj.weight",
            )
        )
        rename_keys.append(
            (
                f"transformer.decoder.layers.{i}.multihead_attn.out_proj.bias",
                f"decoder.layers.{i}.encoder_attn.out_proj.bias",
            )
        )
        rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight"))
        rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias"))
        rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight"))
        rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias"))
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight")
        )
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias")
        )
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight")
        )
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias")
        )
        rename_keys.append(
            (f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight")
        )
        rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias"))

    
    rename_keys.extend(
        [
            ("input_proj.weight", "input_projection.weight"),
            ("input_proj.bias", "input_projection.bias"),
            ("query_embed.weight", "query_position_embeddings.weight"),
            ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
            ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
            ("class_embed.weight", "class_labels_classifier.weight"),
            ("class_embed.bias", "class_labels_classifier.bias"),
            ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
            ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
            ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
            ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
            ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
            ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
        ]
    )

    return rename_keys


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


def read_in_q_k_v(state_dict, is_panoptic=False):
    prefix = ""
    if is_panoptic:
        prefix = "detr."

    
    for i in range(6):
        
        in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias")
        
        state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
        state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
        state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
        state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
        state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
        state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
    
    for i in range(6):
        
        in_proj_weight = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias")
        
        state_dict[f"decoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :]
        state_dict[f"decoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256]
        state_dict[f"decoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :]
        state_dict[f"decoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512]
        state_dict[f"decoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :]
        state_dict[f"decoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:]
        
        in_proj_weight_cross_attn = state_dict.pop(
            f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"
        )
        in_proj_bias_cross_attn = state_dict.pop(f"{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias")
        
        state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.weight"] = in_proj_weight_cross_attn[:256, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.q_proj.bias"] = in_proj_bias_cross_attn[:256]
        state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.weight"] = in_proj_weight_cross_attn[256:512, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.k_proj.bias"] = in_proj_bias_cross_attn[256:512]
        state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.weight"] = in_proj_weight_cross_attn[-256:, :]
        state_dict[f"decoder.layers.{i}.encoder_attn.v_proj.bias"] = in_proj_bias_cross_attn[-256:]



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_detr_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False):
    

    
    config, is_panoptic = get_detr_config(model_name)

    
    model_name_to_original_name = {
        "detr-resnet-50": "detr_resnet50",
        "detr-resnet-101": "detr_resnet101",
    }
    logger.info(f"Converting model {model_name}...")
    detr = torch.hub.load("facebookresearch/detr", model_name_to_original_name[model_name], pretrained=True).eval()
    state_dict = detr.state_dict()
    
    for src, dest in create_rename_keys(config):
        if is_panoptic:
            src = "detr." + src
        rename_key(state_dict, src, dest)
    
    read_in_q_k_v(state_dict, is_panoptic=is_panoptic)
    
    prefix = "detr.model." if is_panoptic else "model."
    for key in state_dict.copy().keys():
        if is_panoptic:
            if (
                key.startswith("detr")
                and not key.startswith("class_labels_classifier")
                and not key.startswith("bbox_predictor")
            ):
                val = state_dict.pop(key)
                state_dict["detr.model" + key[4:]] = val
            elif "class_labels_classifier" in key or "bbox_predictor" in key:
                val = state_dict.pop(key)
                state_dict["detr." + key] = val
            elif key.startswith("bbox_attention") or key.startswith("mask_head"):
                continue
            else:
                val = state_dict.pop(key)
                state_dict[prefix + key] = val
        else:
            if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"):
                val = state_dict.pop(key)
                state_dict[prefix + key] = val

    
    model = DetrForSegmentation(config) if is_panoptic else DetrForObjectDetection(config)
    model.load_state_dict(state_dict)
    model.eval()

    
    format = "coco_panoptic" if is_panoptic else "coco_detection"
    processor = DetrImageProcessor(format=format)

    encoding = processor(images=prepare_img(), return_tensors="pt")
    pixel_values = encoding["pixel_values"]

    original_outputs = detr(pixel_values)
    outputs = model(pixel_values)

    assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-3)
    assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-3)
    if is_panoptic:
        assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4)
    print("Looks ok!")

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

    if push_to_hub:
        
        logger.info("Uploading PyTorch model and image processor to the hub...")
        model.push_to_hub(f"nielsr/{model_name}")
        processor.push_to_hub(f"nielsr/{model_name}")


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

    parser.add_argument(
        "--model_name",
        default="detr-resnet-50",
        type=str,
        choices=["detr-resnet-50", "detr-resnet-101"],
        help="Name of the DETR model you'd like to convert.",
    )
    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 push the model to the hub or not.")
    args = parser.parse_args()
    convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)