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 owlv2 checkpoints from the original repository url https github comgoogleresearchscenictreemainscenicprojectsowlvit import argparse import collections import os import jax import jax numpy as jnp import numpy as np import torch from flax training import checkpoints from huggingfacehub import hfhubdownload from pil import image from transformers import cliptokenizer owlv2config owlv2forobjectdetection owlv2imageprocessor owlv2processor owlv2textconfig owlv2visionconfig from transformers utils import logging logging setverbosityinfo logger logging getloggername def getowlv2configmodelname if large in modelname imagesize 1008 patchsize 14 visionhiddensize 1024 visionintermediatesize 4096 visionnumhiddenlayers 24 visionnumattentionheads 16 projectiondim 768 texthiddensize 768 textintermediatesize 3072 textnumattentionheads 12 textnumhiddenlayers 12 else imagesize 960 patchsize 16 visionhiddensize 768 visionintermediatesize 3072 visionnumhiddenlayers 12 visionnumattentionheads 12 projectiondim 512 texthiddensize 512 textintermediatesize 2048 textnumattentionheads 8 textnumhiddenlayers 12 visionconfig owlv2visionconfig patchsizepatchsize imagesizeimagesize hiddensizevisionhiddensize numhiddenlayersvisionnumhiddenlayers intermediatesizevisionintermediatesize numattentionheadsvisionnumattentionheads textconfig owlv2textconfig hiddensizetexthiddensize intermediatesizetextintermediatesize numattentionheadstextnumattentionheads numhiddenlayerstextnumhiddenlayers config owlv2config textconfigtextconfig todict visionconfigvisionconfig todict projectiondimprojectiondim return config def flattennesteddictparams parentkey sep items for k v in params items newkey parentkey sep k if parentkey else k if isinstancev collections mutablemapping items extendflattennesteddictv newkey sepsep items else items appendnewkey v return dictitems here we list all keys to be renamed original name on the left our name on the right def createrenamekeysconfig modelname renamekeys fmt off clip vision encoder renamekeys appendbackboneclipvisualclassembedding owlv2 visionmodel embeddings classembedding renamekeys appendbackboneclipvisualconv1kernel owlv2 visionmodel embeddings patchembedding weight renamekeys appendbackboneclipvisualpositionalembedding owlv2 visionmodel embeddings positionembedding weight renamekeys appendbackboneclipvisuallnprescale owlv2 visionmodel prelayernorm weight renamekeys appendbackboneclipvisuallnprebias owlv2 visionmodel prelayernorm bias for i in rangeconfig visionconfig numhiddenlayers if v2 in modelname renamekeys appendfbackboneclipvisualtransformerresblocks iln0scale fowlv2 visionmodel encoder layers i layernorm1 weight renamekeys appendfbackboneclipvisualtransformerresblocks iln0bias fowlv2 visionmodel encoder layers i layernorm1 bias renamekeys appendfbackboneclipvisualtransformerresblocks iln1scale fowlv2 visionmodel encoder layers i layernorm2 weight renamekeys appendfbackboneclipvisualtransformerresblocks iln1bias fowlv2 visionmodel encoder layers i layernorm2 bias else renamekeys appendfbackboneclipvisualtransformerresblocks iln1scale fowlv2 visionmodel encoder layers i layernorm1 weight renamekeys appendfbackboneclipvisualtransformerresblocks iln1bias fowlv2 visionmodel encoder layers i layernorm1 bias renamekeys appendfbackboneclipvisualtransformerresblocks iln2scale fowlv2 visionmodel encoder layers i layernorm2 weight renamekeys appendfbackboneclipvisualtransformerresblocks iln2bias fowlv2 visionmodel encoder layers i layernorm2 bias renamekeys appendfbackboneclipvisualtransformerresblocks imlpcfckernel fowlv2 visionmodel encoder layers i mlp fc1 weight renamekeys appendfbackboneclipvisualtransformerresblocks imlpcfcbias fowlv2 visionmodel encoder layers i mlp fc1 bias renamekeys appendfbackboneclipvisualtransformerresblocks imlpcprojkernel fowlv2 visionmodel encoder layers i mlp fc2 weight renamekeys appendfbackboneclipvisualtransformerresblocks imlpcprojbias fowlv2 visionmodel encoder layers i mlp fc2 bias renamekeys appendfbackboneclipvisualtransformerresblocks iattnquerykernel fowlv2 visionmodel encoder layers i selfattn qproj weight renamekeys appendfbackboneclipvisualtransformerresblocks iattnquerybias fowlv2 visionmodel encoder layers i selfattn qproj bias renamekeys appendfbackboneclipvisualtransformerresblocks iattnkeykernel fowlv2 visionmodel encoder layers i selfattn kproj weight renamekeys appendfbackboneclipvisualtransformerresblocks iattnkeybias fowlv2 visionmodel encoder layers i selfattn kproj bias renamekeys appendfbackboneclipvisualtransformerresblocks iattnvaluekernel fowlv2 visionmodel encoder layers i selfattn vproj weight renamekeys appendfbackboneclipvisualtransformerresblocks iattnvaluebias fowlv2 visionmodel encoder layers i selfattn vproj bias renamekeys appendfbackboneclipvisualtransformerresblocks iattnoutkernel fowlv2 visionmodel encoder layers i selfattn outproj weight renamekeys appendfbackboneclipvisualtransformerresblocks iattnoutbias fowlv2 visionmodel encoder layers i selfattn outproj bias renamekeys appendbackboneclipvisuallnpostscale owlv2 visionmodel postlayernorm weight renamekeys appendbackboneclipvisuallnpostbias owlv2 visionmodel postlayernorm bias clip text encoder renamekeys appendbackbonecliptexttokenembeddingembedding owlv2 textmodel embeddings tokenembedding weight renamekeys appendbackbonecliptextpositionalembedding owlv2 textmodel embeddings positionembedding weight for i in rangeconfig textconfig numhiddenlayers if v2 in modelname renamekeys appendfbackbonecliptexttransformerresblocks iln0scale fowlv2 textmodel encoder layers i layernorm1 weight renamekeys appendfbackbonecliptexttransformerresblocks iln0bias fowlv2 textmodel encoder layers i layernorm1 bias renamekeys appendfbackbonecliptexttransformerresblocks iln1scale fowlv2 textmodel encoder layers i layernorm2 weight renamekeys appendfbackbonecliptexttransformerresblocks iln1bias fowlv2 textmodel encoder layers i layernorm2 bias else renamekeys appendfbackbonecliptexttransformerresblocks iln1scale fowlv2 textmodel encoder layers i layernorm1 weight renamekeys appendfbackbonecliptexttransformerresblocks iln1bias fowlv2 textmodel encoder layers i layernorm1 bias renamekeys appendfbackbonecliptexttransformerresblocks iln2scale fowlv2 textmodel encoder layers i layernorm2 weight renamekeys appendfbackbonecliptexttransformerresblocks iln2bias fowlv2 textmodel encoder layers i layernorm2 bias renamekeys appendfbackbonecliptexttransformerresblocks imlpcfckernel fowlv2 textmodel encoder layers i mlp fc1 weight renamekeys appendfbackbonecliptexttransformerresblocks imlpcfcbias fowlv2 textmodel encoder layers i mlp fc1 bias renamekeys appendfbackbonecliptexttransformerresblocks imlpcprojkernel fowlv2 textmodel encoder layers i mlp fc2 weight renamekeys appendfbackbonecliptexttransformerresblocks imlpcprojbias fowlv2 textmodel encoder layers i mlp fc2 bias renamekeys appendfbackbonecliptexttransformerresblocks iattnquerykernel fowlv2 textmodel encoder layers i selfattn qproj weight renamekeys appendfbackbonecliptexttransformerresblocks iattnquerybias fowlv2 textmodel encoder layers i selfattn qproj bias renamekeys appendfbackbonecliptexttransformerresblocks iattnkeykernel fowlv2 textmodel encoder layers i selfattn kproj weight renamekeys appendfbackbonecliptexttransformerresblocks iattnkeybias fowlv2 textmodel encoder layers i selfattn kproj bias renamekeys appendfbackbonecliptexttransformerresblocks iattnvaluekernel fowlv2 textmodel encoder layers i selfattn vproj weight renamekeys appendfbackbonecliptexttransformerresblocks iattnvaluebias fowlv2 textmodel encoder layers i selfattn vproj bias renamekeys appendfbackbonecliptexttransformerresblocks iattnoutkernel fowlv2 textmodel encoder layers i selfattn outproj weight renamekeys appendfbackbonecliptexttransformerresblocks iattnoutbias fowlv2 textmodel encoder layers i selfattn outproj bias renamekeys appendbackbonecliptextlnfinalscale owlv2 textmodel finallayernorm weight renamekeys appendbackbonecliptextlnfinalbias owlv2 textmodel finallayernorm bias logit scale renamekeys appendbackbonecliplogitscale owlv2 logitscale projection heads renamekeys appendbackbonecliptexttextprojectionkernel owlv2 textprojection weight class and box heads renamekeys appendbackbonemergedclasstokenscale layernorm weight renamekeys appendbackbonemergedclasstokenbias layernorm bias renamekeys appendclassheaddense0kernel classhead dense0 weight renamekeys appendclassheaddense0bias classhead dense0 bias renamekeys appendclassheadlogitshiftkernel classhead logitshift weight renamekeys appendclassheadlogitscalekernel classhead logitscale weight renamekeys appendclassheadlogitscalebias classhead logitscale bias renamekeys appendclassheadlogitshiftbias classhead logitshift bias renamekeys appendobjboxheaddense0kernel boxhead dense0 weight renamekeys appendobjboxheaddense0bias boxhead dense0 bias renamekeys appendobjboxheaddense1kernel boxhead dense1 weight renamekeys appendobjboxheaddense1bias boxhead dense1 bias renamekeys appendobjboxheaddense2kernel boxhead dense2 weight renamekeys appendobjboxheaddense2bias boxhead dense2 bias objectness head only for v2 if v2 in modelname renamekeys appendobjectnessheaddense0kernel objectnesshead dense0 weight renamekeys appendobjectnessheaddense0bias objectnesshead dense0 bias renamekeys appendobjectnessheaddense1kernel objectnesshead dense1 weight renamekeys appendobjectnessheaddense1bias objectnesshead dense1 bias renamekeys appendobjectnessheaddense2kernel objectnesshead dense2 weight renamekeys appendobjectnessheaddense2bias objectnesshead dense2 bias fmt on return renamekeys def renameandreshapekeydct old new config val dct popold if outproj in new or vproj in new or kproj in new or qproj in new and vision in new val val reshape1 config visionconfig hiddensize if outproj in new or vproj in new or kproj in new or qproj in new and text in new val val reshape1 config textconfig hiddensize if patchembedding in new printreshaping patch embedding for new val val transpose3 2 0 1 elif new endswithweight and positionembedding not in new and tokenembedding not in new val val t if new endswithbias val val reshape1 dctnew torch fromnumpynp arrayval torch nograd def convertowlv2checkpointmodelname checkpointpath pytorchdumpfolderpath pushtohub verifylogits config getowlv2configmodelname see available checkpoints at https github comgoogleresearchscenictreemainscenicprojectsowlvitpretrainedcheckpoints variables checkpoints restorecheckpointcheckpointpath targetnone variables variablesparams if v2 in modelname else variablesoptimizertarget flaxparams jax treeutil treemaplambda x x astypejnp float32 if x dtype jnp bfloat16 else x variables statedict flattennesteddictflaxparams rename keys renamekeys createrenamekeysconfig modelname for src dest in renamekeys renameandreshapekeystatedict src dest config load huggingface model model owlv2forobjectdetectionconfig missingkeys unexpectedkeys model loadstatedictstatedict strictfalse assert missingkeys owlv2 visualprojection weight assert unexpectedkeys model eval initialize image processor size height config visionconfig imagesize width config visionconfig imagesize imageprocessor owlv2imageprocessorsizesize initialize tokenizer tokenizer cliptokenizer frompretrainedopenaiclipvitbasepatch32 padtoken modelmaxlength16 initialize processor processor owlv2processorimageprocessorimageprocessor tokenizertokenizer verify pixelvalues and inputids filepath hfhubdownloadrepoidnielsrtestimage filenameowlvitpixelvalues960 pt repotypedataset originalpixelvalues torch loadfilepath permute0 3 1 2 filepath hfhubdownloadrepoidnielsrtestimage filenameowlv2inputids pt repotypedataset originalinputids torch loadfilepath squeeze filepath hfhubdownloadrepoidadirikowlvit repotypespace filenameassetsastronaut png image image openfilepath texts face rocket nasa badge starspangled banner inputs processortexttexts imagesimage returntensorspt if large not in modelname assert torch allcloseinputs pixelvalues originalpixelvalues float atol1e6 assert torch allcloseinputs inputids 4 originalinputids 4 atol1e6 with torch nograd outputs modelinputs logits outputs logits predboxes outputs predboxes objectnesslogits outputs objectnesslogits if verifylogits if modelname owlv2basepatch16 expectedlogits torch tensor 10 0043 9 0226 8 0433 12 4569 14 0380 12 6153 21 0731 22 2705 21 8850 expectedboxes torch tensor 0 0136 0 0223 0 0269 0 0406 0 0327 0 0797 0 0638 0 1539 0 1255 expectedobjectnesslogits torch tensor 5 6589 7 7702 16 3965 elif modelname owlv2basepatch16finetuned expectedlogits torch tensor 9 2391 9 2313 8 0295 14 5498 16 8450 14 7166 15 1278 17 3060 15 7169 expectedboxes torch tensor 0 0103 0 0094 0 0207 0 0483 0 0729 0 1013 0 0629 0 1396 0 1313 expectedobjectnesslogits torch tensor 6 5234 13 3788 14 6627 elif modelname owlv2basepatch16ensemble expectedlogits torch tensor 8 6353 9 5409 6 6154 7 9442 9 6151 6 7117 12 4593 15 3332 12 1048 expectedboxes torch tensor 0 0126 0 0090 0 0238 0 0387 0 0227 0 0754 0 0582 0 1058 0 1139 expectedobjectnesslogits torch tensor 6 0628 5 9507 10 4486 elif modelname owlv2largepatch14 expectedlogits torch tensor 12 6662 11 8384 12 1880 16 0599 16 5835 16 9364 21 4957 26 7038 25 1313 expectedboxes torch tensor 0 0136 0 0161 0 0256 0 0126 0 0135 0 0202 0 0498 0 0948 0 0915 expectedobjectnesslogits torch tensor 6 7196 9 4590 13 9472 elif modelname owlv2largepatch14finetuned expectedlogits torch tensor 9 5413 9 7130 7 9762 9 5731 9 7277 8 2252 15 4434 19 3084 16 5490 expectedboxes torch tensor 0 0089 0 0080 0 0175 0 0112 0 0098 0 0179 0 0375 0 0821 0 0528 expectedobjectnesslogits torch tensor 6 2655 6 5845 11 3105 elif modelname owlv2largepatch14ensemble expectedlogits torch tensor 12 2037 12 2070 11 5371 13 4875 13 8235 13 1586 18 2007 22 9834 20 6816 expectedboxes torch tensor 0 0126 0 0127 0 0222 0 0107 0 0113 0 0164 0 0482 0 1162 0 0885 expectedobjectnesslogits torch tensor 7 7572 8 3637 13 0334 printobjectness logits objectnesslogits 3 3 printlogits logits0 3 3 printpred boxes predboxes0 3 3 assert torch allcloselogits0 3 3 expectedlogits atol1e3 assert torch allclosepredboxes0 3 3 expectedboxes atol1e3 assert torch allcloseobjectnesslogits 3 3 expectedobjectnesslogits atol1e3 printlooks ok else printmodel converted without verifying logits if pytorchdumpfolderpath is not none printsaving model and processor locally create folder to save model if not os path isdirpytorchdumpfolderpath os mkdirpytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath processor savepretrainedpytorchdumpfolderpath if pushtohub printfpushing modelname to the hub model pushtohubfgooglemodelname processor pushtohubfgooglemodelname if name main parser argparse argumentparser required parameters parser addargument modelname defaultowlv2basepatch16 choices owlv2basepatch16 owlv2basepatch16finetuned owlv2basepatch16ensemble owlv2largepatch14 owlv2largepatch14finetuned owlv2largepatch14ensemble typestr helpname of the owlv2 model you d like to convert from flax to pytorch parser addargument checkpointpath defaultnone typestr requiredtrue helppath to the original flax checkpoint parser addargument pytorchdumpfolderpath defaultnone typestr requiredfalse helppath to the output pytorch model directory parser addargument verifylogits actionstorefalse requiredfalse helppath to the output pytorch model directory parser addargumentpushtohub actionstoretrue helppush model and image preprocessor to the hub args parser parseargs convertowlv2checkpoint args modelname args checkpointpath args pytorchdumpfolderpath args pushtohub args verifylogits 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 owlv2 checkpoints from the original repository url https github com google research scenic tree main scenic projects owl_vit here we list all keys to be renamed original name on the left our name on the right fmt off clip vision encoder clip text encoder logit scale projection heads class and box heads objectness head only for v2 fmt on copy paste tweak model s weights to our owl vit structure see available checkpoints at https github com google research scenic tree main scenic projects owl_vit pretrained checkpoints rename keys load huggingface model initialize image processor initialize tokenizer initialize processor verify pixel_values and input_ids create folder to save model required parameters
import argparse
import collections
import os

import jax
import jax.numpy as jnp
import numpy as np
import torch
from flax.training import checkpoints
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    CLIPTokenizer,
    Owlv2Config,
    Owlv2ForObjectDetection,
    Owlv2ImageProcessor,
    Owlv2Processor,
    Owlv2TextConfig,
    Owlv2VisionConfig,
)
from transformers.utils import logging


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


def get_owlv2_config(model_name):
    if "large" in model_name:
        image_size = 1008
        patch_size = 14
        vision_hidden_size = 1024
        vision_intermediate_size = 4096
        vision_num_hidden_layers = 24
        vision_num_attention_heads = 16
        projection_dim = 768
        text_hidden_size = 768
        text_intermediate_size = 3072
        text_num_attention_heads = 12
        text_num_hidden_layers = 12
    else:
        image_size = 960
        patch_size = 16
        vision_hidden_size = 768
        vision_intermediate_size = 3072
        vision_num_hidden_layers = 12
        vision_num_attention_heads = 12
        projection_dim = 512
        text_hidden_size = 512
        text_intermediate_size = 2048
        text_num_attention_heads = 8
        text_num_hidden_layers = 12

    vision_config = Owlv2VisionConfig(
        patch_size=patch_size,
        image_size=image_size,
        hidden_size=vision_hidden_size,
        num_hidden_layers=vision_num_hidden_layers,
        intermediate_size=vision_intermediate_size,
        num_attention_heads=vision_num_attention_heads,
    )
    text_config = Owlv2TextConfig(
        hidden_size=text_hidden_size,
        intermediate_size=text_intermediate_size,
        num_attention_heads=text_num_attention_heads,
        num_hidden_layers=text_num_hidden_layers,
    )

    config = Owlv2Config(
        text_config=text_config.to_dict(),
        vision_config=vision_config.to_dict(),
        projection_dim=projection_dim,
    )

    return config


def flatten_nested_dict(params, parent_key="", sep="/"):
    items = []

    for k, v in params.items():
        new_key = parent_key + sep + k if parent_key else k

        if isinstance(v, collections.MutableMapping):
            items.extend(flatten_nested_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)



def create_rename_keys(config, model_name):
    rename_keys = []

    
    
    rename_keys.append(("backbone/clip/visual/class_embedding", "owlv2.vision_model.embeddings.class_embedding"))
    rename_keys.append(("backbone/clip/visual/conv1/kernel", "owlv2.vision_model.embeddings.patch_embedding.weight"))
    rename_keys.append(("backbone/clip/visual/positional_embedding", "owlv2.vision_model.embeddings.position_embedding.weight"))
    rename_keys.append(("backbone/clip/visual/ln_pre/scale", "owlv2.vision_model.pre_layernorm.weight"))
    rename_keys.append(("backbone/clip/visual/ln_pre/bias", "owlv2.vision_model.pre_layernorm.bias"))

    for i in range(config.vision_config.num_hidden_layers):
        if "v2" in model_name:
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_0/scale", f"owlv2.vision_model.encoder.layers.{i}.layer_norm1.weight"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_0/bias", f"owlv2.vision_model.encoder.layers.{i}.layer_norm1.bias"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_1/scale", f"owlv2.vision_model.encoder.layers.{i}.layer_norm2.weight"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_1/bias", f"owlv2.vision_model.encoder.layers.{i}.layer_norm2.bias"))
        else:
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_1/scale", f"owlv2.vision_model.encoder.layers.{i}.layer_norm1.weight"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_1/bias", f"owlv2.vision_model.encoder.layers.{i}.layer_norm1.bias"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_2/scale", f"owlv2.vision_model.encoder.layers.{i}.layer_norm2.weight"))
            rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/ln_2/bias", f"owlv2.vision_model.encoder.layers.{i}.layer_norm2.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/mlp/c_fc/kernel", f"owlv2.vision_model.encoder.layers.{i}.mlp.fc1.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/mlp/c_fc/bias", f"owlv2.vision_model.encoder.layers.{i}.mlp.fc1.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/mlp/c_proj/kernel", f"owlv2.vision_model.encoder.layers.{i}.mlp.fc2.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/mlp/c_proj/bias", f"owlv2.vision_model.encoder.layers.{i}.mlp.fc2.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/query/kernel", f"owlv2.vision_model.encoder.layers.{i}.self_attn.q_proj.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/query/bias", f"owlv2.vision_model.encoder.layers.{i}.self_attn.q_proj.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/key/kernel", f"owlv2.vision_model.encoder.layers.{i}.self_attn.k_proj.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/key/bias", f"owlv2.vision_model.encoder.layers.{i}.self_attn.k_proj.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/value/kernel", f"owlv2.vision_model.encoder.layers.{i}.self_attn.v_proj.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/value/bias", f"owlv2.vision_model.encoder.layers.{i}.self_attn.v_proj.bias"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/out/kernel", f"owlv2.vision_model.encoder.layers.{i}.self_attn.out_proj.weight"))
        rename_keys.append((f"backbone/clip/visual/transformer/resblocks.{i}/attn/out/bias", f"owlv2.vision_model.encoder.layers.{i}.self_attn.out_proj.bias"))

    rename_keys.append(("backbone/clip/visual/ln_post/scale", "owlv2.vision_model.post_layernorm.weight"))
    rename_keys.append(("backbone/clip/visual/ln_post/bias", "owlv2.vision_model.post_layernorm.bias"))

    
    rename_keys.append(("backbone/clip/text/token_embedding/embedding", "owlv2.text_model.embeddings.token_embedding.weight"))
    rename_keys.append(("backbone/clip/text/positional_embedding", "owlv2.text_model.embeddings.position_embedding.weight"))

    for i in range(config.text_config.num_hidden_layers):
        if "v2" in model_name:
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_0/scale", f"owlv2.text_model.encoder.layers.{i}.layer_norm1.weight"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_0/bias", f"owlv2.text_model.encoder.layers.{i}.layer_norm1.bias"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_1/scale", f"owlv2.text_model.encoder.layers.{i}.layer_norm2.weight"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_1/bias", f"owlv2.text_model.encoder.layers.{i}.layer_norm2.bias"))
        else:
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_1/scale", f"owlv2.text_model.encoder.layers.{i}.layer_norm1.weight"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_1/bias", f"owlv2.text_model.encoder.layers.{i}.layer_norm1.bias"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_2/scale", f"owlv2.text_model.encoder.layers.{i}.layer_norm2.weight"))
            rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/ln_2/bias", f"owlv2.text_model.encoder.layers.{i}.layer_norm2.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/mlp/c_fc/kernel", f"owlv2.text_model.encoder.layers.{i}.mlp.fc1.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/mlp/c_fc/bias", f"owlv2.text_model.encoder.layers.{i}.mlp.fc1.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/mlp/c_proj/kernel", f"owlv2.text_model.encoder.layers.{i}.mlp.fc2.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/mlp/c_proj/bias", f"owlv2.text_model.encoder.layers.{i}.mlp.fc2.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/query/kernel", f"owlv2.text_model.encoder.layers.{i}.self_attn.q_proj.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/query/bias", f"owlv2.text_model.encoder.layers.{i}.self_attn.q_proj.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/key/kernel", f"owlv2.text_model.encoder.layers.{i}.self_attn.k_proj.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/key/bias", f"owlv2.text_model.encoder.layers.{i}.self_attn.k_proj.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/value/kernel", f"owlv2.text_model.encoder.layers.{i}.self_attn.v_proj.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/value/bias", f"owlv2.text_model.encoder.layers.{i}.self_attn.v_proj.bias"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/out/kernel", f"owlv2.text_model.encoder.layers.{i}.self_attn.out_proj.weight"))
        rename_keys.append((f"backbone/clip/text/transformer/resblocks.{i}/attn/out/bias", f"owlv2.text_model.encoder.layers.{i}.self_attn.out_proj.bias"))

    rename_keys.append(("backbone/clip/text/ln_final/scale", "owlv2.text_model.final_layer_norm.weight"))
    rename_keys.append(("backbone/clip/text/ln_final/bias", "owlv2.text_model.final_layer_norm.bias"))

    
    rename_keys.append(("backbone/clip/logit_scale", "owlv2.logit_scale"))

    
    rename_keys.append(("backbone/clip/text/text_projection/kernel", "owlv2.text_projection.weight"))

    
    rename_keys.append(("backbone/merged_class_token/scale", "layer_norm.weight"))
    rename_keys.append(("backbone/merged_class_token/bias", "layer_norm.bias"))
    rename_keys.append(("class_head/Dense_0/kernel", "class_head.dense0.weight"))
    rename_keys.append(("class_head/Dense_0/bias", "class_head.dense0.bias"))
    rename_keys.append(("class_head/logit_shift/kernel", "class_head.logit_shift.weight"))
    rename_keys.append(("class_head/logit_scale/kernel", "class_head.logit_scale.weight"))
    rename_keys.append(("class_head/logit_scale/bias", "class_head.logit_scale.bias"))
    rename_keys.append(("class_head/logit_shift/bias", "class_head.logit_shift.bias"))
    rename_keys.append(("obj_box_head/Dense_0/kernel", "box_head.dense0.weight"))
    rename_keys.append(("obj_box_head/Dense_0/bias", "box_head.dense0.bias"))
    rename_keys.append(("obj_box_head/Dense_1/kernel", "box_head.dense1.weight"))
    rename_keys.append(("obj_box_head/Dense_1/bias", "box_head.dense1.bias"))
    rename_keys.append(("obj_box_head/Dense_2/kernel", "box_head.dense2.weight"))
    rename_keys.append(("obj_box_head/Dense_2/bias", "box_head.dense2.bias"))

    
    if "v2" in model_name:
        rename_keys.append(("objectness_head/Dense_0/kernel", "objectness_head.dense0.weight"))
        rename_keys.append(("objectness_head/Dense_0/bias", "objectness_head.dense0.bias"))
        rename_keys.append(("objectness_head/Dense_1/kernel", "objectness_head.dense1.weight"))
        rename_keys.append(("objectness_head/Dense_1/bias", "objectness_head.dense1.bias"))
        rename_keys.append(("objectness_head/Dense_2/kernel", "objectness_head.dense2.weight"))
        rename_keys.append(("objectness_head/Dense_2/bias", "objectness_head.dense2.bias"))

    

    return rename_keys


def rename_and_reshape_key(dct, old, new, config):
    val = dct.pop(old)

    if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "vision" in new:
        val = val.reshape(-1, config.vision_config.hidden_size)
    if ("out_proj" in new or "v_proj" in new or "k_proj" in new or "q_proj" in new) and "text" in new:
        val = val.reshape(-1, config.text_config.hidden_size)

    if "patch_embedding" in new:
        print("Reshaping patch embedding... for", new)
        val = val.transpose(3, 2, 0, 1)
    elif new.endswith("weight") and "position_embedding" not in new and "token_embedding" not in new:
        val = val.T

    if new.endswith("bias"):
        val = val.reshape(-1)

    dct[new] = torch.from_numpy(np.array(val))


@torch.no_grad()
def convert_owlv2_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub, verify_logits):
    
    config = get_owlv2_config(model_name)

    
    variables = checkpoints.restore_checkpoint(checkpoint_path, target=None)
    variables = variables["params"] if "v2" in model_name else variables["optimizer"]["target"]
    flax_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, variables)
    state_dict = flatten_nested_dict(flax_params)

    
    rename_keys = create_rename_keys(config, model_name)
    for src, dest in rename_keys:
        rename_and_reshape_key(state_dict, src, dest, config)

    
    model = Owlv2ForObjectDetection(config)
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    assert missing_keys == ["owlv2.visual_projection.weight"]
    assert unexpected_keys == []
    model.eval()

    
    size = {"height": config.vision_config.image_size, "width": config.vision_config.image_size}
    image_processor = Owlv2ImageProcessor(size=size)
    
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32", pad_token="!", model_max_length=16)
    
    processor = Owlv2Processor(image_processor=image_processor, tokenizer=tokenizer)

    
    filepath = hf_hub_download(repo_id="nielsr/test-image", filename="owlvit_pixel_values_960.pt", repo_type="dataset")
    original_pixel_values = torch.load(filepath).permute(0, 3, 1, 2)

    filepath = hf_hub_download(repo_id="nielsr/test-image", filename="owlv2_input_ids.pt", repo_type="dataset")
    original_input_ids = torch.load(filepath).squeeze()

    filepath = hf_hub_download(repo_id="adirik/OWL-ViT", repo_type="space", filename="assets/astronaut.png")
    image = Image.open(filepath)
    texts = [["face", "rocket", "nasa badge", "star-spangled banner"]]
    inputs = processor(text=texts, images=image, return_tensors="pt")

    if "large" not in model_name:
        assert torch.allclose(inputs.pixel_values, original_pixel_values.float(), atol=1e-6)
    assert torch.allclose(inputs.input_ids[:4, :], original_input_ids[:4, :], atol=1e-6)

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        pred_boxes = outputs.pred_boxes
        objectness_logits = outputs.objectness_logits

    if verify_logits:
        if model_name == "owlv2-base-patch16":
            expected_logits = torch.tensor(
                [[-10.0043, -9.0226, -8.0433], [-12.4569, -14.0380, -12.6153], [-21.0731, -22.2705, -21.8850]]
            )
            expected_boxes = torch.tensor(
                [[0.0136, 0.0223, 0.0269], [0.0406, 0.0327, 0.0797], [0.0638, 0.1539, 0.1255]]
            )
            expected_objectness_logits = torch.tensor(
                [[-5.6589, -7.7702, -16.3965]],
            )
        elif model_name == "owlv2-base-patch16-finetuned":
            expected_logits = torch.tensor(
                [[-9.2391, -9.2313, -8.0295], [-14.5498, -16.8450, -14.7166], [-15.1278, -17.3060, -15.7169]],
            )
            expected_boxes = torch.tensor(
                [[0.0103, 0.0094, 0.0207], [0.0483, 0.0729, 0.1013], [0.0629, 0.1396, 0.1313]]
            )
            expected_objectness_logits = torch.tensor(
                [[-6.5234, -13.3788, -14.6627]],
            )
        elif model_name == "owlv2-base-patch16-ensemble":
            expected_logits = torch.tensor(
                [[-8.6353, -9.5409, -6.6154], [-7.9442, -9.6151, -6.7117], [-12.4593, -15.3332, -12.1048]]
            )
            expected_boxes = torch.tensor(
                [[0.0126, 0.0090, 0.0238], [0.0387, 0.0227, 0.0754], [0.0582, 0.1058, 0.1139]]
            )
            expected_objectness_logits = torch.tensor(
                [[-6.0628, -5.9507, -10.4486]],
            )
        elif model_name == "owlv2-large-patch14":
            expected_logits = torch.tensor(
                [[-12.6662, -11.8384, -12.1880], [-16.0599, -16.5835, -16.9364], [-21.4957, -26.7038, -25.1313]],
            )
            expected_boxes = torch.tensor(
                [[0.0136, 0.0161, 0.0256], [0.0126, 0.0135, 0.0202], [0.0498, 0.0948, 0.0915]],
            )
            expected_objectness_logits = torch.tensor(
                [[-6.7196, -9.4590, -13.9472]],
            )
        elif model_name == "owlv2-large-patch14-finetuned":
            expected_logits = torch.tensor(
                [[-9.5413, -9.7130, -7.9762], [-9.5731, -9.7277, -8.2252], [-15.4434, -19.3084, -16.5490]],
            )
            expected_boxes = torch.tensor(
                [[0.0089, 0.0080, 0.0175], [0.0112, 0.0098, 0.0179], [0.0375, 0.0821, 0.0528]],
            )
            expected_objectness_logits = torch.tensor(
                [[-6.2655, -6.5845, -11.3105]],
            )
        elif model_name == "owlv2-large-patch14-ensemble":
            expected_logits = torch.tensor(
                [[-12.2037, -12.2070, -11.5371], [-13.4875, -13.8235, -13.1586], [-18.2007, -22.9834, -20.6816]],
            )
            expected_boxes = torch.tensor(
                [[0.0126, 0.0127, 0.0222], [0.0107, 0.0113, 0.0164], [0.0482, 0.1162, 0.0885]],
            )
            expected_objectness_logits = torch.tensor(
                [[-7.7572, -8.3637, -13.0334]],
            )

        print("Objectness logits:", objectness_logits[:3, :3])
        print("Logits:", logits[0, :3, :3])
        print("Pred boxes:", pred_boxes[0, :3, :3])

        assert torch.allclose(logits[0, :3, :3], expected_logits, atol=1e-3)
        assert torch.allclose(pred_boxes[0, :3, :3], expected_boxes, atol=1e-3)
        assert torch.allclose(objectness_logits[:3, :3], expected_objectness_logits, atol=1e-3)
        print("Looks ok!")
    else:
        print("Model converted without verifying logits")

    if pytorch_dump_folder_path is not None:
        print("Saving model and processor locally...")
        
        if not os.path.isdir(pytorch_dump_folder_path):
            os.mkdir(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_name} to the hub...")
        model.push_to_hub(f"google/{model_name}")
        processor.push_to_hub(f"google/{model_name}")


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

    
    parser.add_argument(
        "--model_name",
        default="owlv2-base-patch16",
        choices=[
            "owlv2-base-patch16",
            "owlv2-base-patch16-finetuned",
            "owlv2-base-patch16-ensemble",
            "owlv2-large-patch14",
            "owlv2-large-patch14-finetuned",
            "owlv2-large-patch14-ensemble",
        ],
        type=str,
        help="Name of the Owlv2 model you'd like to convert from FLAX to PyTorch.",
    )
    parser.add_argument(
        "--checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the original Flax checkpoint.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=False,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument(
        "--verify_logits",
        action="store_false",
        required=False,
        help="Path to the output PyTorch model directory.",
    )
    parser.add_argument("--push_to_hub", action="store_true", help="Push model and image preprocessor to the hub")

    args = parser.parse_args()
    convert_owlv2_checkpoint(
        args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub, args.verify_logits
    )