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 git checkpoints from the original repository url https github commicrosoftgenerativeimage2texttreemain import argparse from pathlib import path import numpy as np import requests import torch from huggingfacehub import hfhubdownload from pil import image from torchvision transforms import centercrop compose normalize resize totensor from transformers import autotokenizer clipimageprocessor gitconfig gitforcausallm gitprocessor gitvisionconfig videomaeimageprocessor from transformers utils import logging logging setverbosityinfo logger logging getloggername def getgitconfigmodelname if base in modelname and vqa in modelname imagesize 480 elif large in modelname and vqa in modelname imagesize 420 else imagesize 224 visionconfig gitvisionconfigimagesizeimagesize if large in modelname visionconfig patchsize 14 visionconfig hiddensize 1024 visionconfig intermediatesize 4096 visionconfig numhiddenlayers 24 visionconfig numattentionheads 16 isvideo vatex in modelname or msrvtt in modelname numimagewithembedding 6 if isvideo else none config gitconfigvisionconfigvisionconfig todict numimagewithembeddingnumimagewithembedding return config imagesize isvideo here we list all keys to be renamed original name on the left our name on the right def createrenamekeysconfig prefix renamekeys image encoder ftm off renamekeys append fprefiximageencoder classembedding git imageencoder visionmodel embeddings classembedding renamekeys append fprefiximageencoder positionalembedding git imageencoder visionmodel embeddings positionembedding weight renamekeys append fprefiximageencoder conv1 weight git imageencoder visionmodel embeddings patchembedding weight renamekeys appendfprefiximageencoder lnpre weight git imageencoder visionmodel prelayrnorm weight renamekeys appendfprefiximageencoder lnpre bias git imageencoder visionmodel prelayrnorm bias renamekeys append fprefiximageencoder lnpost weight git imageencoder visionmodel postlayernorm weight renamekeys appendfprefiximageencoder lnpost bias git imageencoder visionmodel postlayernorm bias fmt on renamekeys appendfprefiximageencoder proj git imageencoder visualprojection weight fmt off for i in rangeconfig visionconfig numhiddenlayers image encoder layers output projection 2 feedforward neural networks and 2 layernorms renamekeys appendfprefiximageencoder transformer resblocks i attn outproj weight fgit imageencoder visionmodel encoder layers i selfattn outproj weight renamekeys appendfprefiximageencoder transformer resblocks i attn outproj bias fgit imageencoder visionmodel encoder layers i selfattn outproj bias renamekeys appendfprefiximageencoder transformer resblocks i ln1 weight fgit imageencoder visionmodel encoder layers i layernorm1 weight renamekeys appendfprefiximageencoder transformer resblocks i ln1 bias fgit imageencoder visionmodel encoder layers i layernorm1 bias renamekeys appendfprefiximageencoder transformer resblocks i mlp cfc weight fgit imageencoder visionmodel encoder layers i mlp fc1 weight renamekeys appendfprefiximageencoder transformer resblocks i mlp cfc bias fgit imageencoder visionmodel encoder layers i mlp fc1 bias renamekeys appendfprefiximageencoder transformer resblocks i mlp cproj weight fgit imageencoder visionmodel encoder layers i mlp fc2 weight renamekeys appendfprefiximageencoder transformer resblocks i mlp cproj bias fgit imageencoder visionmodel encoder layers i mlp fc2 bias renamekeys appendfprefiximageencoder transformer resblocks i ln2 weight fgit imageencoder visionmodel encoder layers i layernorm2 weight renamekeys appendfprefiximageencoder transformer resblocks i ln2 bias fgit imageencoder visionmodel encoder layers i layernorm2 bias fmt on text decoder fmt off renamekeys appendfprefixtextual embedding words weight git embeddings wordembeddings weight renamekeys appendfprefixtextual embedding positions weight git embeddings positionembeddings weight renamekeys appendfprefixtextual visualprojection 0 weight git visualprojection visualprojection 0 weight renamekeys appendfprefixtextual visualprojection 0 bias git visualprojection visualprojection 0 bias renamekeys appendfprefixtextual visualprojection 1 weight git visualprojection visualprojection 1 weight renamekeys appendfprefixtextual visualprojection 1 bias git visualprojection visualprojection 1 bias renamekeys appendfprefixtextual embedding layernorm weight git embeddings layernorm weight renamekeys appendfprefixtextual embedding layernorm bias git embeddings layernorm bias renamekeys appendfprefixtextual output weight output weight renamekeys appendfprefixtextual output bias output bias for i in rangeconfig numhiddenlayers renamekeys appendfprefixtextual transformer encoder layer i attention self query weight fgit encoder layer i attention self query weight renamekeys appendfprefixtextual transformer encoder layer i attention self query bias fgit encoder layer i attention self query bias renamekeys appendfprefixtextual transformer encoder layer i attention self key weight fgit encoder layer i attention self key weight renamekeys appendfprefixtextual transformer encoder layer i attention self key bias fgit encoder layer i attention self key bias renamekeys appendfprefixtextual transformer encoder layer i attention self value weight fgit encoder layer i attention self value weight renamekeys appendfprefixtextual transformer encoder layer i attention self value bias fgit encoder layer i attention self value bias renamekeys appendfprefixtextual transformer encoder layer i attention output dense weight fgit encoder layer i attention output dense weight renamekeys appendfprefixtextual transformer encoder layer i attention output dense bias fgit encoder layer i attention output dense bias renamekeys appendfprefixtextual transformer encoder layer i attention output layernorm weight fgit encoder layer i attention output layernorm weight renamekeys appendfprefixtextual transformer encoder layer i attention output layernorm bias fgit encoder layer i attention output layernorm bias renamekeys appendfprefixtextual transformer encoder layer i intermediate dense weight fgit encoder layer i intermediate dense weight renamekeys appendfprefixtextual transformer encoder layer i intermediate dense bias fgit encoder layer i intermediate dense bias renamekeys appendfprefixtextual transformer encoder layer i output dense weight fgit encoder layer i output dense weight renamekeys appendfprefixtextual transformer encoder layer i output dense bias fgit encoder layer i output dense bias renamekeys appendfprefixtextual transformer encoder layer i output layernorm weight fgit encoder layer i output layernorm weight renamekeys appendfprefixtextual transformer encoder layer i output layernorm bias fgit encoder layer i output layernorm bias fmt on if config numimagewithembedding is not none renamekeys appendimgtemperalembedding 0 git imgtemperalembedding 0 renamekeys appendimgtemperalembedding 1 git imgtemperalembedding 1 renamekeys appendimgtemperalembedding 2 git imgtemperalembedding 2 renamekeys appendimgtemperalembedding 3 git imgtemperalembedding 3 renamekeys appendimgtemperalembedding 4 git imgtemperalembedding 4 renamekeys appendimgtemperalembedding 5 git imgtemperalembedding 5 return renamekeys def renamekeydct old new val dct popold dctnew val t if imageencoder visualprojection in new else val we split up the matrix of each clip encoder layer into queries keys and values def readinqkvstatedict config prefix dim config visionconfig hiddensize for i in rangeconfig visionconfig numhiddenlayers read in weights bias of input projection layer in the original implementation this is a single matrix bias inprojweight statedict popfprefiximageencoder transformer resblocks i attn inprojweight inprojbias statedict popfprefiximageencoder transformer resblocks i attn inprojbias next add query keys and values in that order to the state dict statedictfgit imageencoder visionmodel encoder layers i selfattn qproj weight inprojweight dim statedictfgit imageencoder visionmodel encoder layers i selfattn qproj bias inprojbias dim statedictfgit imageencoder visionmodel encoder layers i selfattn kproj weight inprojweight dim dim 2 statedictfgit imageencoder visionmodel encoder layers i selfattn kproj bias inprojbias dim dim 2 statedictfgit imageencoder visionmodel encoder layers i selfattn vproj weight inprojweight dim statedictfgit imageencoder visionmodel encoder layers i selfattn vproj bias inprojbiasdim we will verify our results on an image def prepareimgmodelname if textvqa in modelname filepath hfhubdownloadrepoidnielsrtextvqasample filenamebus png repotypedataset image image openfilepath convertrgb else url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw return image def preparevideo from decord import videoreader cpu set seed for reproducability np random seed0 def sampleframeindicescliplen framesamplerate seglen convertedlen intcliplen framesamplerate endidx np random randintconvertedlen seglen startidx endidx convertedlen indices np linspacestartidx endidx numcliplen indices np clipindices startidx endidx 1 astypenp int64 return indices video clip consists of 300 frames 10 seconds at 30 fps filepath hfhubdownloadrepoidnielsrvideodemo filenameeatingspaghetti mp4 repotypedataset videoreader videoreaderfilepath numthreads1 ctxcpu0 sample 6 frames videoreader seek0 indices sampleframeindicescliplen6 framesamplerate4 seglenlenvideoreader video videoreader getbatchindices asnumpy return video torch nograd def convertgitcheckpointmodelname pytorchdumpfolderpath pushtohubfalse modelnametourl gitbase https publicgit blob core windows netdataoutputgitbasesnapshotmodel pt gitbasecoco https publicgit blob core windows netdataoutputgitbasecocosnapshotmodel pt gitbasetextcaps https publicgit blob core windows netdataoutputgitbasetextcapssnapshotmodel pt gitbasevqav2 https publicgit blob core windows netdataoutputgitbasevqav2snapshotmodel pt gitbasetextvqa https publicgit blob core windows netdataoutputgitbasetextvqasnapshotmodel pt todo gitbasevatex https publicgit blob core windows netdataoutputgitbasevatexsnapshotmodel pt gitbasemsrvttqa https publicgit blob core windows netdataoutputgitbasemsrvttqasnapshotmodel pt gitlarge https publicgit blob core windows netdataoutputgitlargesnapshotmodel pt gitlargecoco https publicgit blob core windows netdataoutputgitlargecocosnapshotmodel pt gitlargetextcaps https publicgit blob core windows netdataoutputgitlargetextcapssnapshotmodel pt gitlargevqav2 https publicgit blob core windows netdataoutputgitlargevqav2snapshotmodel pt gitlargetextvqa https publicgit blob core windows netdataoutputgitlargetextvqasnapshotmodel pt gitlargevatex https publicgit blob core windows netdataoutputgitlargevatexsnapshotmodel pt gitlargemsrvttqa https publicgit blob core windows netdataoutputgitlargemsrvttqasnapshotmodel pt gitlarger https publicgit blob core windows netdataoutputgitlargersnapshotmodel pt gitlargercoco https publicgit blob core windows netdataoutputgitlargercocosnapshotmodel pt gitlargertextcaps https publicgit blob core windows netdataoutputgitlargertextcapssnapshotmodel pt modelnametopath gitlarge usersnielsroggedocumentsgitgitlargemodel pt gitlargecoco usersnielsroggedocumentsgitgitlargecocomodel pt gitlargetextcaps usersnielsroggedocumentsgitgitlargetextcapsmodel pt gitlargevqav2 usersnielsroggedocumentsgitgitlargevqav2model pt gitlargetextvqa usersnielsroggedocumentsgitgitlargetextvqamodel pt define git configuration based on model name config imagesize isvideo getgitconfigmodelname if large in modelname and not isvideo and larger not in modelname large checkpoints take way too long to download checkpointpath modelnametopathmodelname statedict torch loadcheckpointpath maplocationcpumodel else checkpointurl modelnametourlmodelname statedict torch hub loadstatedictfromurlcheckpointurl maplocationcpu filenamemodelname model rename keys prefix module if modelname gitbase else renamekeys createrenamekeysconfig prefixprefix for src dest in renamekeys renamekeystatedict src dest readinqkvstatedict config prefixprefix load huggingface model model gitforcausallmconfig missingkeys unexpectedkeys model loadstatedictstatedict strictfalse model eval printmissing keys missingkeys printunexpected keys unexpectedkeys assert missingkeys git embeddings positionids git imageencoder visionmodel embeddings positionids assert unexpectedkeys git imageencoder visualprojection weight verify results imageprocessor videomaeimageprocessor sizeshortestedge imagesize cropsizeheight imagesize width imagesize if isvideo else clipimageprocessor sizeshortestedge imagesize cropsizeheight imagesize width imagesize tokenizer autotokenizer frompretrainedbertbaseuncased modelinputnamesinputids attentionmask processor gitprocessortokenizertokenizer imageprocessorimageprocessor if isvideo video preparevideo pixelvalues processorimageslistvideo returntensorspt pixelvalues else image prepareimgmodelname imagetransforms compose resizeimagesize interpolationimage bicubic centercropimagesize totensor normalize0 48145466 0 4578275 0 40821073 0 26862954 0 26130258 0 27577711 originalpixelvalues imagetransformsimage unsqueeze0 pixelvalues processorimagesimage returntensorspt pixelvalues assert torch allclosepixelvalues originalpixelvalues inputids torch tensor101 outputs modelinputids pixelvaluespixelvalues logits outputs logits printlogits logits0 1 3 if modelname gitbase expectedslicelogits torch tensor1 2832 1 2835 1 2840 elif modelname gitbasecoco expectedslicelogits torch tensor0 9925 0 9930 0 9935 elif modelname gitbasetextcaps expectedslicelogits torch tensor1 2980 1 2983 1 2985 elif modelname gitbasevqav2 expectedslicelogits torch tensor0 8570 0 8568 0 8561 elif modelname gitbasetextvqa expectedslicelogits torch tensor1 4085 1 4083 1 4082 elif modelname gitbasevatex expectedslicelogits torch tensor1 3451 1 3447 1 3447 elif modelname gitbasemsrvttqa expectedslicelogits torch tensor0 8554 0 8550 0 8540 elif modelname gitlarge expectedslicelogits torch tensor1 1708 1 1707 1 1705 elif modelname gitlargecoco expectedslicelogits torch tensor1 0425 1 0423 1 0422 elif modelname gitlargetextcaps expectedslicelogits torch tensor1 2705 1 2708 1 2706 elif modelname gitlargevqav2 expectedslicelogits torch tensor0 7042 0 7043 0 7043 elif modelname gitlargetextvqa expectedslicelogits torch tensor0 8590 0 8592 0 8590 elif modelname gitlargevatex expectedslicelogits torch tensor1 0113 1 0114 1 0113 elif modelname gitlargemsrvttqa expectedslicelogits torch tensor0 0130 0 0134 0 0131 elif modelname gitlarger expectedslicelogits torch tensor1 1283 1 1285 1 1286 elif modelname gitlargercoco expectedslicelogits torch tensor0 9641 0 9641 0 9641 elif modelname gitlargertextcaps expectedslicelogits torch tensor1 1121 1 1120 1 1124 assert torch allcloselogits0 1 3 expectedslicelogits atol1e4 printlooks ok prompt if textvqa in modelname prompt what does the front of the bus say at the top elif msrvttqa in modelname prompt what does the woman eat elif vqa in modelname prompt what are the cats doing inputids tokenizerprompt addspecialtokensfalse inputids inputids processor tokenizer clstokenid inputids inputids torch tensorinputids unsqueeze0 printgenerating caption generatedids model generatepixelvaluespixelvalues inputidsinputids maxlength50 printgenerated caption processor batchdecodegeneratedids skipspecialtokenstrue if pytorchdumpfolderpath is not none pathpytorchdumpfolderpath mkdirexistoktrue printfsaving model and processor of modelname to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub printfpushing model and processor of modelname to the hub model pushtohubfmicrosoftmodelname processor pushtohubfmicrosoftmodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultgitbase typestr 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 args parser parseargs convertgitcheckpointargs modelname args pytorchdumpfolderpath args pushtohub 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 git checkpoints from the original repository url https github com microsoft generativeimage2text tree main here we list all keys to be renamed original name on the left our name on the right image encoder ftm off fmt on fmt off image encoder layers output projection 2 feedforward neural networks and 2 layernorms fmt on text decoder fmt off fmt on we split up the matrix of each clip encoder layer into queries keys and values read in weights bias of input projection layer in the 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 set seed for reproducability sample a given number of frame indices from the video args clip_len int total number of frames to sample frame_sample_rate int sample every n th frame seg_len int maximum allowed index of sample s last frame returns indices list int list of sampled frame indices video clip consists of 300 frames 10 seconds at 30 fps sample 6 frames copy paste tweak model s weights to our git structure todo define git configuration based on model name large checkpoints take way too long to download rename keys load huggingface model verify results required parameters
import argparse
from pathlib import Path

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor

from transformers import (
    AutoTokenizer,
    CLIPImageProcessor,
    GitConfig,
    GitForCausalLM,
    GitProcessor,
    GitVisionConfig,
    VideoMAEImageProcessor,
)
from transformers.utils import logging


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


def get_git_config(model_name):
    if "base" in model_name and "vqa" in model_name:
        image_size = 480
    elif "large" in model_name and "vqa" in model_name:
        image_size = 420
    else:
        image_size = 224

    vision_config = GitVisionConfig(image_size=image_size)

    if "large" in model_name:
        vision_config.patch_size = 14
        vision_config.hidden_size = 1024
        vision_config.intermediate_size = 4096
        vision_config.num_hidden_layers = 24
        vision_config.num_attention_heads = 16

    is_video = "vatex" in model_name or "msrvtt" in model_name
    num_image_with_embedding = 6 if is_video else None
    config = GitConfig(vision_config=vision_config.to_dict(), num_image_with_embedding=num_image_with_embedding)

    return config, image_size, is_video



def create_rename_keys(config, prefix=""):
    rename_keys = []

    
    
    rename_keys.append(
        (f"{prefix}image_encoder.class_embedding", "git.image_encoder.vision_model.embeddings.class_embedding")
    )
    rename_keys.append(
        (
            f"{prefix}image_encoder.positional_embedding",
            "git.image_encoder.vision_model.embeddings.position_embedding.weight",
        )
    )
    rename_keys.append(
        (f"{prefix}image_encoder.conv1.weight", "git.image_encoder.vision_model.embeddings.patch_embedding.weight")
    )
    rename_keys.append((f"{prefix}image_encoder.ln_pre.weight", "git.image_encoder.vision_model.pre_layrnorm.weight"))
    rename_keys.append((f"{prefix}image_encoder.ln_pre.bias", "git.image_encoder.vision_model.pre_layrnorm.bias"))
    rename_keys.append(
        (f"{prefix}image_encoder.ln_post.weight", "git.image_encoder.vision_model.post_layernorm.weight")
    )
    rename_keys.append((f"{prefix}image_encoder.ln_post.bias", "git.image_encoder.vision_model.post_layernorm.bias"))
    
    rename_keys.append((f"{prefix}image_encoder.proj", "git.image_encoder.visual_projection.weight"))

    
    for i in range(config.vision_config.num_hidden_layers):
        
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.attn.out_proj.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.out_proj.weight"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.attn.out_proj.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.out_proj.bias"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_1.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm1.weight"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_1.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm1.bias"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_fc.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc1.weight"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_fc.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc1.bias"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_proj.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc2.weight"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.mlp.c_proj.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.mlp.fc2.bias"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_2.weight", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm2.weight"))
        rename_keys.append((f"{prefix}image_encoder.transformer.resblocks.{i}.ln_2.bias", f"git.image_encoder.vision_model.encoder.layers.{i}.layer_norm2.bias"))
    

    
    
    rename_keys.append((f"{prefix}textual.embedding.words.weight", "git.embeddings.word_embeddings.weight"))
    rename_keys.append((f"{prefix}textual.embedding.positions.weight", "git.embeddings.position_embeddings.weight"))
    rename_keys.append((f"{prefix}textual.visual_projection.0.weight", "git.visual_projection.visual_projection.0.weight"))
    rename_keys.append((f"{prefix}textual.visual_projection.0.bias", "git.visual_projection.visual_projection.0.bias"))
    rename_keys.append((f"{prefix}textual.visual_projection.1.weight", "git.visual_projection.visual_projection.1.weight"))
    rename_keys.append((f"{prefix}textual.visual_projection.1.bias", "git.visual_projection.visual_projection.1.bias"))

    rename_keys.append((f"{prefix}textual.embedding.layer_norm.weight", "git.embeddings.LayerNorm.weight"))
    rename_keys.append((f"{prefix}textual.embedding.layer_norm.bias", "git.embeddings.LayerNorm.bias"))
    rename_keys.append((f"{prefix}textual.output.weight", "output.weight"))
    rename_keys.append((f"{prefix}textual.output.bias", "output.bias"))
    for i in range(config.num_hidden_layers):
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.query.weight", f"git.encoder.layer.{i}.attention.self.query.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.query.bias", f"git.encoder.layer.{i}.attention.self.query.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.key.weight", f"git.encoder.layer.{i}.attention.self.key.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.key.bias", f"git.encoder.layer.{i}.attention.self.key.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.value.weight", f"git.encoder.layer.{i}.attention.self.value.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.self.value.bias", f"git.encoder.layer.{i}.attention.self.value.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.dense.weight", f"git.encoder.layer.{i}.attention.output.dense.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.dense.bias", f"git.encoder.layer.{i}.attention.output.dense.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.LayerNorm.weight", f"git.encoder.layer.{i}.attention.output.LayerNorm.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.attention.output.LayerNorm.bias", f"git.encoder.layer.{i}.attention.output.LayerNorm.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.intermediate.dense.weight", f"git.encoder.layer.{i}.intermediate.dense.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.intermediate.dense.bias", f"git.encoder.layer.{i}.intermediate.dense.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.dense.weight", f"git.encoder.layer.{i}.output.dense.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.dense.bias", f"git.encoder.layer.{i}.output.dense.bias"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.LayerNorm.weight", f"git.encoder.layer.{i}.output.LayerNorm.weight"))
        rename_keys.append((f"{prefix}textual.transformer.encoder.layer.{i}.output.LayerNorm.bias", f"git.encoder.layer.{i}.output.LayerNorm.bias"))
    

    if config.num_image_with_embedding is not None:
        rename_keys.append(("img_temperal_embedding.0", "git.img_temperal_embedding.0"))
        rename_keys.append(("img_temperal_embedding.1", "git.img_temperal_embedding.1"))
        rename_keys.append(("img_temperal_embedding.2", "git.img_temperal_embedding.2"))
        rename_keys.append(("img_temperal_embedding.3", "git.img_temperal_embedding.3"))
        rename_keys.append(("img_temperal_embedding.4", "git.img_temperal_embedding.4"))
        rename_keys.append(("img_temperal_embedding.5", "git.img_temperal_embedding.5"))

    return rename_keys


def rename_key(dct, old, new):
    val = dct.pop(old)
    dct[new] = val.T if "image_encoder.visual_projection" in new else val



def read_in_q_k_v(state_dict, config, prefix=""):
    dim = config.vision_config.hidden_size
    for i in range(config.vision_config.num_hidden_layers):
        
        in_proj_weight = state_dict.pop(f"{prefix}image_encoder.transformer.resblocks.{i}.attn.in_proj_weight")
        in_proj_bias = state_dict.pop(f"{prefix}image_encoder.transformer.resblocks.{i}.attn.in_proj_bias")
        
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[
            :dim, :
        ]
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:dim]
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[
            dim : dim * 2, :
        ]
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[
            dim : dim * 2
        ]
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[
            -dim:, :
        ]
        state_dict[f"git.image_encoder.vision_model.encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-dim:]



def prepare_img(model_name):
    if "textvqa" in model_name:
        filepath = hf_hub_download(repo_id="nielsr/textvqa-sample", filename="bus.png", repo_type="dataset")
        image = Image.open(filepath).convert("RGB")
    else:
        url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        image = Image.open(requests.get(url, stream=True).raw)

    return image


def prepare_video():
    from decord import VideoReader, cpu

    
    np.random.seed(0)

    def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
        
        converted_len = int(clip_len * frame_sample_rate)
        end_idx = np.random.randint(converted_len, seg_len)
        start_idx = end_idx - converted_len
        indices = np.linspace(start_idx, end_idx, num=clip_len)
        indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
        return indices

    
    file_path = hf_hub_download(repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset")
    videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))

    
    videoreader.seek(0)
    indices = sample_frame_indices(clip_len=6, frame_sample_rate=4, seg_len=len(videoreader))
    video = videoreader.get_batch(indices).asnumpy()

    return video


@torch.no_grad()
def convert_git_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
    

    model_name_to_url = {
        "git-base": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE/snapshot/model.pt",
        "git-base-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_COCO/snapshot/model.pt",
        "git-base-textcaps": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_TEXTCAPS/snapshot/model.pt",
        "git-base-vqav2": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_VQAv2/snapshot/model.pt",
        "git-base-textvqa": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_TEXTVQA/snapshot/model.pt",  
        "git-base-vatex": "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_VATEX/snapshot/model.pt",
        "git-base-msrvtt-qa": (
            "https://publicgit.blob.core.windows.net/data/output/GIT_BASE_MSRVTT_QA/snapshot/model.pt"
        ),
        "git-large": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE/snapshot/model.pt",
        "git-large-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_COCO/snapshot/model.pt",
        "git-large-textcaps": (
            "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_TEXTCAPS/snapshot/model.pt"
        ),
        "git-large-vqav2": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_VQAv2/snapshot/model.pt",
        "git-large-textvqa": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_TEXTVQA/snapshot/model.pt",
        "git-large-vatex": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_VATEX/snapshot/model.pt",
        "git-large-msrvtt-qa": (
            "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_MSRVTT_QA/snapshot/model.pt"
        ),
        "git-large-r": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R/snapshot/model.pt",
        "git-large-r-coco": "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R_COCO/snapshot/model.pt",
        "git-large-r-textcaps": (
            "https://publicgit.blob.core.windows.net/data/output/GIT_LARGE_R_TEXTCAPS/snapshot/model.pt"
        ),
    }

    model_name_to_path = {
        "git-large": "/Users/nielsrogge/Documents/GIT/git_large_model.pt",
        "git-large-coco": "/Users/nielsrogge/Documents/GIT/git_large_coco_model.pt",
        "git-large-textcaps": "/Users/nielsrogge/Documents/GIT/git_large_textcaps_model.pt",
        "git-large-vqav2": "/Users/nielsrogge/Documents/GIT/git_large_vqav2_model.pt",
        "git-large-textvqa": "/Users/nielsrogge/Documents/GIT/git_large_textvqa_model.pt",
    }

    
    config, image_size, is_video = get_git_config(model_name)
    if "large" in model_name and not is_video and "large-r" not in model_name:
        
        checkpoint_path = model_name_to_path[model_name]
        state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
    else:
        checkpoint_url = model_name_to_url[model_name]
        state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu", file_name=model_name)[
            "model"
        ]
    
    prefix = "module." if model_name == "git-base" else ""
    rename_keys = create_rename_keys(config, prefix=prefix)
    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, prefix=prefix)

    
    model = GitForCausalLM(config)
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    model.eval()

    print("Missing keys:", missing_keys)
    print("Unexpected keys:", unexpected_keys)

    assert missing_keys == ["git.embeddings.position_ids", "git.image_encoder.vision_model.embeddings.position_ids"]
    assert unexpected_keys == ["git.image_encoder.visual_projection.weight"]

    
    image_processor = (
        VideoMAEImageProcessor(
            size={"shortest_edge": image_size}, crop_size={"height": image_size, "width": image_size}
        )
        if is_video
        else CLIPImageProcessor(
            size={"shortest_edge": image_size}, crop_size={"height": image_size, "width": image_size}
        )
    )
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", model_input_names=["input_ids", "attention_mask"])
    processor = GitProcessor(tokenizer=tokenizer, image_processor=image_processor)

    if is_video:
        video = prepare_video()
        pixel_values = processor(images=list(video), return_tensors="pt").pixel_values
    else:
        image = prepare_img(model_name)
        image_transforms = Compose(
            [
                Resize(image_size, interpolation=Image.BICUBIC),
                CenterCrop(image_size),
                ToTensor(),
                Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
            ]
        )
        original_pixel_values = image_transforms(image).unsqueeze(0)
        pixel_values = processor(images=image, return_tensors="pt").pixel_values

        assert torch.allclose(pixel_values, original_pixel_values)

    input_ids = torch.tensor([[101]])
    outputs = model(input_ids, pixel_values=pixel_values)
    logits = outputs.logits
    print("Logits:", logits[0, -1, :3])

    if model_name == "git-base":
        expected_slice_logits = torch.tensor([-1.2832, -1.2835, -1.2840])
    elif model_name == "git-base-coco":
        expected_slice_logits = torch.tensor([-0.9925, -0.9930, -0.9935])
    elif model_name == "git-base-textcaps":
        expected_slice_logits = torch.tensor([-1.2980, -1.2983, -1.2985])
    elif model_name == "git-base-vqav2":
        expected_slice_logits = torch.tensor([-0.8570, -0.8568, -0.8561])
    elif model_name == "git-base-textvqa":
        expected_slice_logits = torch.tensor([-1.4085, -1.4083, -1.4082])
    elif model_name == "git-base-vatex":
        expected_slice_logits = torch.tensor([-1.3451, -1.3447, -1.3447])
    elif model_name == "git-base-msrvtt-qa":
        expected_slice_logits = torch.tensor([-0.8554, -0.8550, -0.8540])
    elif model_name == "git-large":
        expected_slice_logits = torch.tensor([-1.1708, -1.1707, -1.1705])
    elif model_name == "git-large-coco":
        expected_slice_logits = torch.tensor([-1.0425, -1.0423, -1.0422])
    elif model_name == "git-large-textcaps":
        expected_slice_logits = torch.tensor([-1.2705, -1.2708, -1.2706])
    elif model_name == "git-large-vqav2":
        expected_slice_logits = torch.tensor([-0.7042, -0.7043, -0.7043])
    elif model_name == "git-large-textvqa":
        expected_slice_logits = torch.tensor([-0.8590, -0.8592, -0.8590])
    elif model_name == "git-large-vatex":
        expected_slice_logits = torch.tensor([-1.0113, -1.0114, -1.0113])
    elif model_name == "git-large-msrvtt-qa":
        expected_slice_logits = torch.tensor([0.0130, 0.0134, 0.0131])
    elif model_name == "git-large-r":
        expected_slice_logits = torch.tensor([-1.1283, -1.1285, -1.1286])
    elif model_name == "git-large-r-coco":
        expected_slice_logits = torch.tensor([-0.9641, -0.9641, -0.9641])
    elif model_name == "git-large-r-textcaps":
        expected_slice_logits = torch.tensor([-1.1121, -1.1120, -1.1124])

    assert torch.allclose(logits[0, -1, :3], expected_slice_logits, atol=1e-4)
    print("Looks ok!")

    prompt = ""
    if "textvqa" in model_name:
        prompt = "what does the front of the bus say at the top?"
    elif "msrvtt-qa" in model_name:
        prompt = "what does the woman eat?"
    elif "vqa" in model_name:
        prompt = "what are the cats doing?"
    input_ids = tokenizer(prompt, add_special_tokens=False).input_ids
    input_ids = [processor.tokenizer.cls_token_id] + input_ids
    input_ids = torch.tensor(input_ids).unsqueeze(0)
    print("Generating caption...")
    generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
    print("Generated caption:", processor.batch_decode(generated_ids, skip_special_tokens=True))

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

    if push_to_hub:
        print(f"Pushing model and processor of {model_name} to the hub...")
        model.push_to_hub(f"microsoft/{model_name}")
        processor.push_to_hub(f"microsoft/{model_name}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--model_name",
        default="git-base",
        type=str,
        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.",
    )

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