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 owlvit checkpoints from the original repository url https github comgoogleresearchscenictreemainscenicprojectsowlvit import argparse import collections import jax import jax numpy as jnp import torch import torch nn as nn from clip model import clip from flax training import checkpoints from huggingfacehub import repository from transformers import cliptokenizer owlvitconfig owlvitforobjectdetection owlvitimageprocessor owlvitmodel owlvitprocessor configs vitb32 embeddim 512 imageresolution 768 contextlength 16 vocabsize 49408 visionlayers 12 visionwidth 768 visionpatchsize 32 transformerwidth 512 transformerheads 8 transformerlayers 12 vitb16 embeddim 512 imageresolution 768 contextlength 16 vocabsize 49408 visionlayers 12 visionwidth 768 visionpatchsize 16 transformerwidth 512 transformerheads 8 transformerlayers 12 vitl14 embeddim 768 imageresolution 840 contextlength 16 vocabsize 49408 visionlayers 24 visionwidth 1024 visionpatchsize 14 transformerwidth 768 transformerheads 12 transformerlayers 12 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 def tof32params return jax treeutil treemaplambda x x astypejnp float32 if x dtype jnp bfloat16 else x params def copyattnlayerhfattnlayer ptattnlayer qproj kproj vproj ptattnlayer inprojweight chunk3 dim0 qprojbias kprojbias vprojbias ptattnlayer inprojbias chunk3 dim0 outprojweights ptattnlayer outproj weight outprojbias ptattnlayer outproj bias hfattnlayer qproj weight data qproj hfattnlayer qproj bias data qprojbias hfattnlayer kproj weight data kproj hfattnlayer kproj bias data kprojbias hfattnlayer vproj weight data vproj hfattnlayer vproj bias data vprojbias hfattnlayer outproj weight outprojweights hfattnlayer outproj bias outprojbias def copymlphfmlp ptmlp copylinearhfmlp fc1 ptmlp cfc copylinearhfmlp fc2 ptmlp cproj def copylinearhflinear ptlinear hflinear weight ptlinear weight hflinear bias ptlinear bias def copylayerhflayer ptlayer copy layer norms copylinearhflayer layernorm1 ptlayer ln1 copylinearhflayer layernorm2 ptlayer ln2 copy mlp copymlphflayer mlp ptlayer mlp copy attn copyattnlayerhflayer selfattn ptlayer attn def copylayershflayers ptlayers for hflayer ptlayer in ziphflayers ptlayers copylayerhflayer ptlayer def copyencoderhfencoder ptmodel copy embeds hfencoder embeddings tokenembedding weight ptmodel tokenembedding weight hfencoder embeddings positionembedding weight data ptmodel positionalembedding copy layer norm copylinearhfencoder finallayernorm ptmodel lnfinal copy hidden layers copylayershfencoder encoder layers ptmodel transformer resblocks def copytextmodelandprojectionhfmodel ptmodel copy projection hfmodel textprojection weight data ptmodel textprojection data t copy text encoder copyencoderhfmodel textmodel ptmodel def copyvisionmodelandprojectionhfmodel ptmodel copy projection hfmodel visualprojection weight data ptmodel visual proj data t copy layer norms copylinearhfmodel visionmodel prelayernorm ptmodel visual lnpre copylinearhfmodel visionmodel postlayernorm ptmodel visual lnpost copy embeds hfmodel visionmodel embeddings patchembedding weight data ptmodel visual conv1 weight data hfmodel visionmodel embeddings classembedding ptmodel visual classembedding hfmodel visionmodel embeddings positionembedding weight data ptmodel visual positionalembedding data copy encoder copylayershfmodel visionmodel encoder layers ptmodel visual transformer resblocks def copyclassmergetokenhfmodel flaxparams flaxclasstokenparams flattennesteddictflaxparamsbackbonemergedclasstoken weight torch fromnumpyflaxclasstokenparamsscale bias torch fromnumpyflaxclasstokenparamsbias hfmodel layernorm weight nn parameterweight hfmodel layernorm bias nn parameterbias def copyclassboxheadshfmodel flaxparams ptparams hfmodel statedict newparams rename class prediction head flax params to pytorch hf flaxclassparams flattennesteddictflaxparamsclasshead for flaxkey v in flaxclassparams items torchkey flaxkey replace torchkey torchkey replace kernel weight torchkey torchkey replacedense0 dense0 torchkey classhead torchkey if weight in torchkey and v ndim 2 v v t newparamstorchkey nn parametertorch fromnumpyv rename box prediction box flax params to pytorch hf flaxboxparams flattennesteddictflaxparamsobjboxhead for flaxkey v in flaxboxparams items torchkey flaxkey replace torchkey torchkey replace kernel weight torchkey torchkey replace lower torchkey boxhead torchkey if weight in torchkey and v ndim 2 v v t newparamstorchkey nn parametertorch fromnumpyv copy flax params to pytorch params for name param in newparams items if name in ptparams keys ptparamsname copyparam def copyflaxattnparamshfbackbone flaxattnparams for k v in flaxattnparams items if k startswithtransformer torchkey k replacetransformer resblocks textmodel encoder layers else torchkey k replacevisual transformer resblocks visionmodel encoder layers torchkey torchkey replaceattn selfattn torchkey torchkey replacekey kproj torchkey torchkey replacevalue vproj torchkey torchkey replacequery qproj torchkey torchkey replaceout outproj if bias in torchkey and v ndim 2 shape v shape0 v shape1 v v reshapeshape if weight in torchkey and out in torchkey shape v shape0 v shape1 v shape2 v v reshapeshape t if weight in torchkey and out not in torchkey shape v shape0 v shape1 v shape2 v v reshapeshape t copy flax clip attn params to hf pytorch params v torch fromnumpyv hfbackbone statedicttorchkey copyv def convertattnlayersparams newparams processedattnlayers for k v in params items if attn in k base k k rindexattn 5 if base in processedattnlayers continue processedattnlayers appendbase dim paramsbase out weight shape1 newparamsbase outproj weight paramsbase out weight reshapedim dim t newparamsbase outproj bias paramsbase out bias else newparamsk v return newparams def convertclipbackboneflaxparams torchconfig torchmodel cliptorchconfig torchmodel eval torchclipparams torchmodel statedict flaxclipparams flattennesteddictflaxparamsbackboneclip newtorchparams for flaxkey v in flaxclipparams items torchkey flaxkey replace torchkey torchkey replacetext tokenembedding embedding tokenembedding kernel if torchkey startswithtext transformer or torchkey startswithtext textprojection or torchkey startswithtext lnfinal or torchkey startswithtext positionalembedding torchkey torchkey5 torchkey torchkey replacetextprojection kernel textprojection torchkey torchkey replacevisual proj kernel visual proj torchkey torchkey replace scale weight torchkey torchkey replace kernel weight if conv in torchkey or downsample 0 weight in torchkey v v transpose3 2 0 1 elif weight in torchkey and v ndim 2 and embedding not in torchkey fully connected layers are transposed embeddings are not v v t newtorchparamstorchkey v attnparams convertattnlayersnewtorchparams newtorchparams updateattnparams attnparams copy flax clip backbone params to pytorch params for name param in newtorchparams items if name in torchclipparams keys newparam torch fromnumpynewtorchparamsname torchclipparamsname copynewparam else attnparamsname param return torchclipparams torchmodel attnparams torch nograd def convertowlvitcheckpointptbackbone flaxparams attnparams pytorchdumpfolderpath configpathnone repo repositorypytorchdumpfolderpath clonefromfgooglepytorchdumpfolderpath repo gitpull if configpath is not none config owlvitconfig frompretrainedconfigpath else config owlvitconfig hfbackbone owlvitmodelconfig eval hfmodel owlvitforobjectdetectionconfig eval copytextmodelandprojectionhfbackbone ptbackbone copyvisionmodelandprojectionhfbackbone ptbackbone hfbackbone logitscale ptbackbone logitscale copyflaxattnparamshfbackbone attnparams hfmodel owlvit hfbackbone copyclassmergetokenhfmodel flaxparams copyclassboxheadshfmodel flaxparams save hf model hfmodel savepretrainedrepo localdir initialize image processor imageprocessor owlvitimageprocessor sizeconfig visionconfig imagesize cropsizeconfig visionconfig imagesize initialize tokenizer tokenizer cliptokenizer frompretrainedopenaiclipvitbasepatch32 padtoken modelmaxlength16 initialize processor processor owlvitprocessorimageprocessorimageprocessor tokenizertokenizer imageprocessor savepretrainedrepo localdir processor savepretrainedrepo localdir repo gitadd repo gitcommitupload model and processor repo gitpush if name main parser argparse argumentparser required parameters parser addargument owlvitversion defaultnone typestr requiredtrue helpowlvit model name clipb16 clipb32 clipl14 parser addargument owlvitcheckpoint defaultnone typestr requiredtrue helppath to flax model checkpoint parser addargumenthfconfig defaultnone typestr requiredtrue helppath to hf model config parser addargument pytorchdumpfolderpath defaulthfmodel typestr helppath to the output pytorch model args parser parseargs initialize pytoch clip model modelname args owlvitversion if modelname clipb16 torchconfig configsvitb16 elif modelname clipb32 torchconfig configsvitb32 elif modelname clipl14 torchconfig configsvitl14 load from checkpoint and convert params to float32 variables checkpoints restorecheckpointargs owlvitcheckpoint targetnoneoptimizertarget flaxparams jax treeutil treemaplambda x x astypejnp float32 if x dtype jnp bfloat16 else x variables del variables convert clip backbone ptbackboneparams clippt attnparams convertclipbackboneflaxparams torchconfig convertowlvitcheckpointclippt flaxparams attnparams args pytorchdumpfolderpath args hfconfig 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 owl vit checkpoints from the original repository url https github com google research scenic tree main scenic projects owl_vit copy layer norms copy mlp copy attn copy embeds copy layer norm copy hidden layers copy projection copy text encoder copy projection copy layer norms copy embeds copy encoder rename class prediction head flax params to pytorch hf rename box prediction box flax params to pytorch hf copy flax params to pytorch params copy flax clip attn params to hf pytorch params fully connected layers are transposed embeddings are not copy flax clip backbone params to pytorch params copy paste tweak model s weights to transformers design save hf model initialize image processor initialize tokenizer initialize processor required parameters initialize pytoch clip model load from checkpoint and convert params to float 32 convert clip backbone
import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


CONFIGS = {
    "vit_b32": {
        "embed_dim": 512,
        "image_resolution": 768,
        "context_length": 16,
        "vocab_size": 49408,
        "vision_layers": 12,
        "vision_width": 768,
        "vision_patch_size": 32,
        "transformer_width": 512,
        "transformer_heads": 8,
        "transformer_layers": 12,
    },
    "vit_b16": {
        "embed_dim": 512,
        "image_resolution": 768,
        "context_length": 16,
        "vocab_size": 49408,
        "vision_layers": 12,
        "vision_width": 768,
        "vision_patch_size": 16,
        "transformer_width": 512,
        "transformer_heads": 8,
        "transformer_layers": 12,
    },
    "vit_l14": {
        "embed_dim": 768,
        "image_resolution": 840,
        "context_length": 16,
        "vocab_size": 49408,
        "vision_layers": 24,
        "vision_width": 1024,
        "vision_patch_size": 14,
        "transformer_width": 768,
        "transformer_heads": 12,
        "transformer_layers": 12,
    },
}


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 to_f32(params):
    return jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, params)


def copy_attn_layer(hf_attn_layer, pt_attn_layer):
    q_proj, k_proj, v_proj = pt_attn_layer.in_proj_weight.chunk(3, dim=0)
    q_proj_bias, k_proj_bias, v_proj_bias = pt_attn_layer.in_proj_bias.chunk(3, dim=0)

    out_proj_weights = pt_attn_layer.out_proj.weight
    out_proj_bias = pt_attn_layer.out_proj.bias

    hf_attn_layer.q_proj.weight.data = q_proj
    hf_attn_layer.q_proj.bias.data = q_proj_bias

    hf_attn_layer.k_proj.weight.data = k_proj
    hf_attn_layer.k_proj.bias.data = k_proj_bias

    hf_attn_layer.v_proj.weight.data = v_proj
    hf_attn_layer.v_proj.bias.data = v_proj_bias

    hf_attn_layer.out_proj.weight = out_proj_weights
    hf_attn_layer.out_proj.bias = out_proj_bias


def copy_mlp(hf_mlp, pt_mlp):
    copy_linear(hf_mlp.fc1, pt_mlp.c_fc)
    copy_linear(hf_mlp.fc2, pt_mlp.c_proj)


def copy_linear(hf_linear, pt_linear):
    hf_linear.weight = pt_linear.weight
    hf_linear.bias = pt_linear.bias


def copy_layer(hf_layer, pt_layer):
    
    copy_linear(hf_layer.layer_norm1, pt_layer.ln_1)
    copy_linear(hf_layer.layer_norm2, pt_layer.ln_2)

    
    copy_mlp(hf_layer.mlp, pt_layer.mlp)

    
    copy_attn_layer(hf_layer.self_attn, pt_layer.attn)


def copy_layers(hf_layers, pt_layers):
    for hf_layer, pt_layer in zip(hf_layers, pt_layers):
        copy_layer(hf_layer, pt_layer)


def copy_encoder(hf_encoder, pt_model):
    
    hf_encoder.embeddings.token_embedding.weight = pt_model.token_embedding.weight
    hf_encoder.embeddings.position_embedding.weight.data = pt_model.positional_embedding

    
    copy_linear(hf_encoder.final_layer_norm, pt_model.ln_final)

    
    copy_layers(hf_encoder.encoder.layers, pt_model.transformer.resblocks)


def copy_text_model_and_projection(hf_model, pt_model):
    
    hf_model.text_projection.weight.data = pt_model.text_projection.data.T

    
    copy_encoder(hf_model.text_model, pt_model)


def copy_vision_model_and_projection(hf_model, pt_model):
    
    hf_model.visual_projection.weight.data = pt_model.visual.proj.data.T

    
    copy_linear(hf_model.vision_model.pre_layernorm, pt_model.visual.ln_pre)
    copy_linear(hf_model.vision_model.post_layernorm, pt_model.visual.ln_post)

    
    hf_model.vision_model.embeddings.patch_embedding.weight.data = pt_model.visual.conv1.weight.data
    hf_model.vision_model.embeddings.class_embedding = pt_model.visual.class_embedding
    hf_model.vision_model.embeddings.position_embedding.weight.data = pt_model.visual.positional_embedding.data

    
    copy_layers(hf_model.vision_model.encoder.layers, pt_model.visual.transformer.resblocks)


def copy_class_merge_token(hf_model, flax_params):
    flax_class_token_params = flatten_nested_dict(flax_params["backbone"]["merged_class_token"])

    weight = torch.from_numpy(flax_class_token_params["scale"])
    bias = torch.from_numpy(flax_class_token_params["bias"])
    hf_model.layer_norm.weight = nn.Parameter(weight)
    hf_model.layer_norm.bias = nn.Parameter(bias)


def copy_class_box_heads(hf_model, flax_params):
    pt_params = hf_model.state_dict()
    new_params = {}

    
    flax_class_params = flatten_nested_dict(flax_params["class_head"])

    for flax_key, v in flax_class_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = torch_key.replace(".kernel", ".weight")
        torch_key = torch_key.replace("Dense_0", "dense0")
        torch_key = "class_head." + torch_key

        if "weight" in torch_key and v.ndim == 2:
            v = v.T

        new_params[torch_key] = nn.Parameter(torch.from_numpy(v))

    
    flax_box_params = flatten_nested_dict(flax_params["obj_box_head"])

    for flax_key, v in flax_box_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = torch_key.replace(".kernel", ".weight")
        torch_key = torch_key.replace("_", "").lower()
        torch_key = "box_head." + torch_key

        if "weight" in torch_key and v.ndim == 2:
            v = v.T

        new_params[torch_key] = nn.Parameter(torch.from_numpy(v))

    
    for name, param in new_params.items():
        if name in pt_params.keys():
            pt_params[name].copy_(param)


def copy_flax_attn_params(hf_backbone, flax_attn_params):
    for k, v in flax_attn_params.items():
        if k.startswith("transformer"):
            torch_key = k.replace("transformer.resblocks", "text_model.encoder.layers")
        else:
            torch_key = k.replace("visual.transformer.resblocks", "vision_model.encoder.layers")

        torch_key = torch_key.replace("attn", "self_attn")
        torch_key = torch_key.replace("key", "k_proj")
        torch_key = torch_key.replace("value", "v_proj")
        torch_key = torch_key.replace("query", "q_proj")
        torch_key = torch_key.replace("out", "out_proj")

        if "bias" in torch_key and v.ndim == 2:
            shape = v.shape[0] * v.shape[1]
            v = v.reshape(shape)

        if "weight" in torch_key and "out" in torch_key:
            shape = (v.shape[0] * v.shape[1], v.shape[2])
            v = v.reshape(shape).T

        if "weight" in torch_key and "out" not in torch_key:
            shape = (v.shape[0], v.shape[1] * v.shape[2])
            v = v.reshape(shape).T

        
        v = torch.from_numpy(v)
        hf_backbone.state_dict()[torch_key].copy_(v)


def _convert_attn_layers(params):
    new_params = {}
    processed_attn_layers = []

    for k, v in params.items():
        if "attn." in k:
            base = k[: k.rindex("attn.") + 5]
            if base in processed_attn_layers:
                continue

            processed_attn_layers.append(base)
            dim = params[base + "out.weight"].shape[-1]
            new_params[base + "out_proj.weight"] = params[base + "out.weight"].reshape(dim, dim).T
            new_params[base + "out_proj.bias"] = params[base + "out.bias"]
        else:
            new_params[k] = v
    return new_params


def convert_clip_backbone(flax_params, torch_config):
    torch_model = CLIP(**torch_config)
    torch_model.eval()
    torch_clip_params = torch_model.state_dict()

    flax_clip_params = flatten_nested_dict(flax_params["backbone"]["clip"])
    new_torch_params = {}

    for flax_key, v in flax_clip_params.items():
        torch_key = flax_key.replace("/", ".")
        torch_key = torch_key.replace("text.token_embedding.embedding", "token_embedding.kernel")

        if (
            torch_key.startswith("text.transformer")
            or torch_key.startswith("text.text_projection")
            or torch_key.startswith("text.ln_final")
            or torch_key.startswith("text.positional_embedding")
        ):
            torch_key = torch_key[5:]

        torch_key = torch_key.replace("text_projection.kernel", "text_projection")
        torch_key = torch_key.replace("visual.proj.kernel", "visual.proj")
        torch_key = torch_key.replace(".scale", ".weight")
        torch_key = torch_key.replace(".kernel", ".weight")

        if "conv" in torch_key or "downsample.0.weight" in torch_key:
            v = v.transpose(3, 2, 0, 1)

        elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
            
            v = v.T

        new_torch_params[torch_key] = v

    attn_params = _convert_attn_layers(new_torch_params)
    new_torch_params.update(attn_params)
    attn_params = {}

    
    for name, param in new_torch_params.items():
        if name in torch_clip_params.keys():
            new_param = torch.from_numpy(new_torch_params[name])
            torch_clip_params[name].copy_(new_param)
        else:
            attn_params[name] = param

    return torch_clip_params, torch_model, attn_params


@torch.no_grad()
def convert_owlvit_checkpoint(pt_backbone, flax_params, attn_params, pytorch_dump_folder_path, config_path=None):
    
    repo = Repository(pytorch_dump_folder_path, clone_from=f"google/{pytorch_dump_folder_path}")
    repo.git_pull()

    if config_path is not None:
        config = OwlViTConfig.from_pretrained(config_path)
    else:
        config = OwlViTConfig()

    hf_backbone = OwlViTModel(config).eval()
    hf_model = OwlViTForObjectDetection(config).eval()

    copy_text_model_and_projection(hf_backbone, pt_backbone)
    copy_vision_model_and_projection(hf_backbone, pt_backbone)
    hf_backbone.logit_scale = pt_backbone.logit_scale
    copy_flax_attn_params(hf_backbone, attn_params)

    hf_model.owlvit = hf_backbone
    copy_class_merge_token(hf_model, flax_params)
    copy_class_box_heads(hf_model, flax_params)

    
    hf_model.save_pretrained(repo.local_dir)

    
    image_processor = OwlViTImageProcessor(
        size=config.vision_config.image_size, crop_size=config.vision_config.image_size
    )
    
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32", pad_token="!", model_max_length=16)

    
    processor = OwlViTProcessor(image_processor=image_processor, tokenizer=tokenizer)
    image_processor.save_pretrained(repo.local_dir)
    processor.save_pretrained(repo.local_dir)

    repo.git_add()
    repo.git_commit("Upload model and processor")
    repo.git_push()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--owlvit_version",
        default=None,
        type=str,
        required=True,
        help="OWL-ViT model name [clip_b16, clip_b32, clip_l14].",
    )
    parser.add_argument(
        "--owlvit_checkpoint", default=None, type=str, required=True, help="Path to flax model checkpoint."
    )
    parser.add_argument("--hf_config", default=None, type=str, required=True, help="Path to HF model config.")
    parser.add_argument(
        "--pytorch_dump_folder_path", default="hf_model", type=str, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()

    
    model_name = args.owlvit_version
    if model_name == "clip_b16":
        torch_config = CONFIGS["vit_b16"]
    elif model_name == "clip_b32":
        torch_config = CONFIGS["vit_b32"]
    elif model_name == "clip_l14":
        torch_config = CONFIGS["vit_l14"]

    
    variables = checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)["optimizer"]["target"]
    flax_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, variables)
    del variables

    
    pt_backbone_params, clip_pt, attn_params = convert_clip_backbone(flax_params, torch_config)

    convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)