convert_blipdiffusion_to_diffusers.py

"""
This script requires you to build `LAVIS` from source, since the pip version doesn't have BLIP Diffusion. Follow instructions here: https://github.com/salesforce/LAVIS/tree/main.
"""

import argparse
import os
import tempfile

import torch
from lavis.models import load_model_and_preprocess
from transformers import CLIPTokenizer
from transformers.models.blip_2.configuration_blip_2 import Blip2Config

from diffusers import (
    AutoencoderKL,
    PNDMScheduler,
    UNet2DConditionModel,
)
from diffusers.pipelines import BlipDiffusionPipeline
from diffusers.pipelines.blip_diffusion.blip_image_processing import BlipImageProcessor
from diffusers.pipelines.blip_diffusion.modeling_blip2 import Blip2QFormerModel
from diffusers.pipelines.blip_diffusion.modeling_ctx_clip import ContextCLIPTextModel


BLIP2_CONFIG = {
    "vision_config": {
        "hidden_size": 1024,
        "num_hidden_layers": 23,
        "num_attention_heads": 16,
        "image_size": 224,
        "patch_size": 14,
        "intermediate_size": 4096,
        "hidden_act": "quick_gelu",
    },
    "qformer_config": {
        "cross_attention_frequency": 1,
        "encoder_hidden_size": 1024,
        "vocab_size": 30523,
    },
    "num_query_tokens": 16,
}
blip2config = Blip2Config(**BLIP2_CONFIG)


def qformer_model_from_original_config():
    qformer = Blip2QFormerModel(blip2config)
    return qformer


def embeddings_from_original_checkpoint(model, diffuser_embeddings_prefix, original_embeddings_prefix):
    embeddings = {}
    embeddings.update(
        {
            f"{diffuser_embeddings_prefix}.word_embeddings.weight": model[
                f"{original_embeddings_prefix}.word_embeddings.weight"
            ]
        }
    )
    embeddings.update(
        {
            f"{diffuser_embeddings_prefix}.position_embeddings.weight": model[
                f"{original_embeddings_prefix}.position_embeddings.weight"
            ]
        }
    )
    embeddings.update(
        {f"{diffuser_embeddings_prefix}.LayerNorm.weight": model[f"{original_embeddings_prefix}.LayerNorm.weight"]}
    )
    embeddings.update(
        {f"{diffuser_embeddings_prefix}.LayerNorm.bias": model[f"{original_embeddings_prefix}.LayerNorm.bias"]}
    )
    return embeddings


def proj_layer_from_original_checkpoint(model, diffuser_proj_prefix, original_proj_prefix):
    proj_layer = {}
    proj_layer.update({f"{diffuser_proj_prefix}.dense1.weight": model[f"{original_proj_prefix}.dense1.weight"]})
    proj_layer.update({f"{diffuser_proj_prefix}.dense1.bias": model[f"{original_proj_prefix}.dense1.bias"]})
    proj_layer.update({f"{diffuser_proj_prefix}.dense2.weight": model[f"{original_proj_prefix}.dense2.weight"]})
    proj_layer.update({f"{diffuser_proj_prefix}.dense2.bias": model[f"{original_proj_prefix}.dense2.bias"]})
    proj_layer.update({f"{diffuser_proj_prefix}.LayerNorm.weight": model[f"{original_proj_prefix}.LayerNorm.weight"]})
    proj_layer.update({f"{diffuser_proj_prefix}.LayerNorm.bias": model[f"{original_proj_prefix}.LayerNorm.bias"]})
    return proj_layer


def attention_from_original_checkpoint(model, diffuser_attention_prefix, original_attention_prefix):
    attention = {}
    attention.update(
        {
            f"{diffuser_attention_prefix}.attention.query.weight": model[
                f"{original_attention_prefix}.self.query.weight"
            ]
        }
    )
    attention.update(
        {f"{diffuser_attention_prefix}.attention.query.bias": model[f"{original_attention_prefix}.self.query.bias"]}
    )
    attention.update(
        {f"{diffuser_attention_prefix}.attention.key.weight": model[f"{original_attention_prefix}.self.key.weight"]}
    )
    attention.update(
        {f"{diffuser_attention_prefix}.attention.key.bias": model[f"{original_attention_prefix}.self.key.bias"]}
    )
    attention.update(
        {
            f"{diffuser_attention_prefix}.attention.value.weight": model[
                f"{original_attention_prefix}.self.value.weight"
            ]
        }
    )
    attention.update(
        {f"{diffuser_attention_prefix}.attention.value.bias": model[f"{original_attention_prefix}.self.value.bias"]}
    )
    attention.update(
        {f"{diffuser_attention_prefix}.output.dense.weight": model[f"{original_attention_prefix}.output.dense.weight"]}
    )
    attention.update(
        {f"{diffuser_attention_prefix}.output.dense.bias": model[f"{original_attention_prefix}.output.dense.bias"]}
    )
    attention.update(
        {
            f"{diffuser_attention_prefix}.output.LayerNorm.weight": model[
                f"{original_attention_prefix}.output.LayerNorm.weight"
            ]
        }
    )
    attention.update(
        {
            f"{diffuser_attention_prefix}.output.LayerNorm.bias": model[
                f"{original_attention_prefix}.output.LayerNorm.bias"
            ]
        }
    )
    return attention


def output_layers_from_original_checkpoint(model, diffuser_output_prefix, original_output_prefix):
    output_layers = {}
    output_layers.update({f"{diffuser_output_prefix}.dense.weight": model[f"{original_output_prefix}.dense.weight"]})
    output_layers.update({f"{diffuser_output_prefix}.dense.bias": model[f"{original_output_prefix}.dense.bias"]})
    output_layers.update(
        {f"{diffuser_output_prefix}.LayerNorm.weight": model[f"{original_output_prefix}.LayerNorm.weight"]}
    )
    output_layers.update(
        {f"{diffuser_output_prefix}.LayerNorm.bias": model[f"{original_output_prefix}.LayerNorm.bias"]}
    )
    return output_layers


def encoder_from_original_checkpoint(model, diffuser_encoder_prefix, original_encoder_prefix):
    encoder = {}
    for i in range(blip2config.qformer_config.num_hidden_layers):
        encoder.update(
            attention_from_original_checkpoint(
                model, f"{diffuser_encoder_prefix}.{i}.attention", f"{original_encoder_prefix}.{i}.attention"
            )
        )
        encoder.update(
            attention_from_original_checkpoint(
                model, f"{diffuser_encoder_prefix}.{i}.crossattention", f"{original_encoder_prefix}.{i}.crossattention"
            )
        )

        encoder.update(
            {
                f"{diffuser_encoder_prefix}.{i}.intermediate.dense.weight": model[
                    f"{original_encoder_prefix}.{i}.intermediate.dense.weight"
                ]
            }
        )
        encoder.update(
            {
                f"{diffuser_encoder_prefix}.{i}.intermediate.dense.bias": model[
                    f"{original_encoder_prefix}.{i}.intermediate.dense.bias"
                ]
            }
        )
        encoder.update(
            {
                f"{diffuser_encoder_prefix}.{i}.intermediate_query.dense.weight": model[
                    f"{original_encoder_prefix}.{i}.intermediate_query.dense.weight"
                ]
            }
        )
        encoder.update(
            {
                f"{diffuser_encoder_prefix}.{i}.intermediate_query.dense.bias": model[
                    f"{original_encoder_prefix}.{i}.intermediate_query.dense.bias"
                ]
            }
        )

        encoder.update(
            output_layers_from_original_checkpoint(
                model, f"{diffuser_encoder_prefix}.{i}.output", f"{original_encoder_prefix}.{i}.output"
            )
        )
        encoder.update(
            output_layers_from_original_checkpoint(
                model, f"{diffuser_encoder_prefix}.{i}.output_query", f"{original_encoder_prefix}.{i}.output_query"
            )
        )
    return encoder


def visual_encoder_layer_from_original_checkpoint(model, diffuser_prefix, original_prefix):
    visual_encoder_layer = {}

    visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm1.weight": model[f"{original_prefix}.ln_1.weight"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm1.bias": model[f"{original_prefix}.ln_1.bias"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm2.weight": model[f"{original_prefix}.ln_2.weight"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm2.bias": model[f"{original_prefix}.ln_2.bias"]})
    visual_encoder_layer.update(
        {f"{diffuser_prefix}.self_attn.qkv.weight": model[f"{original_prefix}.attn.in_proj_weight"]}
    )
    visual_encoder_layer.update(
        {f"{diffuser_prefix}.self_attn.qkv.bias": model[f"{original_prefix}.attn.in_proj_bias"]}
    )
    visual_encoder_layer.update(
        {f"{diffuser_prefix}.self_attn.projection.weight": model[f"{original_prefix}.attn.out_proj.weight"]}
    )
    visual_encoder_layer.update(
        {f"{diffuser_prefix}.self_attn.projection.bias": model[f"{original_prefix}.attn.out_proj.bias"]}
    )
    visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc1.weight": model[f"{original_prefix}.mlp.c_fc.weight"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc1.bias": model[f"{original_prefix}.mlp.c_fc.bias"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc2.weight": model[f"{original_prefix}.mlp.c_proj.weight"]})
    visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc2.bias": model[f"{original_prefix}.mlp.c_proj.bias"]})

    return visual_encoder_layer


def visual_encoder_from_original_checkpoint(model, diffuser_prefix, original_prefix):
    visual_encoder = {}

    visual_encoder.update(
        {
            f"{diffuser_prefix}.embeddings.class_embedding": model[f"{original_prefix}.class_embedding"]
            .unsqueeze(0)
            .unsqueeze(0)
        }
    )
    visual_encoder.update(
        {
            f"{diffuser_prefix}.embeddings.position_embedding": model[
                f"{original_prefix}.positional_embedding"
            ].unsqueeze(0)
        }
    )
    visual_encoder.update(
        {f"{diffuser_prefix}.embeddings.patch_embedding.weight": model[f"{original_prefix}.conv1.weight"]}
    )
    visual_encoder.update({f"{diffuser_prefix}.pre_layernorm.weight": model[f"{original_prefix}.ln_pre.weight"]})
    visual_encoder.update({f"{diffuser_prefix}.pre_layernorm.bias": model[f"{original_prefix}.ln_pre.bias"]})

    for i in range(blip2config.vision_config.num_hidden_layers):
        visual_encoder.update(
            visual_encoder_layer_from_original_checkpoint(
                model, f"{diffuser_prefix}.encoder.layers.{i}", f"{original_prefix}.transformer.resblocks.{i}"
            )
        )

    visual_encoder.update({f"{diffuser_prefix}.post_layernorm.weight": model["blip.ln_vision.weight"]})
    visual_encoder.update({f"{diffuser_prefix}.post_layernorm.bias": model["blip.ln_vision.bias"]})

    return visual_encoder


def qformer_original_checkpoint_to_diffusers_checkpoint(model):
    qformer_checkpoint = {}
    qformer_checkpoint.update(embeddings_from_original_checkpoint(model, "embeddings", "blip.Qformer.bert.embeddings"))
    qformer_checkpoint.update({"query_tokens": model["blip.query_tokens"]})
    qformer_checkpoint.update(proj_layer_from_original_checkpoint(model, "proj_layer", "proj_layer"))
    qformer_checkpoint.update(
        encoder_from_original_checkpoint(model, "encoder.layer", "blip.Qformer.bert.encoder.layer")
    )
    qformer_checkpoint.update(visual_encoder_from_original_checkpoint(model, "visual_encoder", "blip.visual_encoder"))
    return qformer_checkpoint


def get_qformer(model):
    print("loading qformer")

    qformer = qformer_model_from_original_config()
    qformer_diffusers_checkpoint = qformer_original_checkpoint_to_diffusers_checkpoint(model)

    load_checkpoint_to_model(qformer_diffusers_checkpoint, qformer)

    print("done loading qformer")
    return qformer


def load_checkpoint_to_model(checkpoint, model):
    with tempfile.NamedTemporaryFile(delete=False) as file:
        torch.save(checkpoint, file.name)
        del checkpoint
        model.load_state_dict(torch.load(file.name), strict=False)

    os.remove(file.name)


def save_blip_diffusion_model(model, args):
    qformer = get_qformer(model)
    qformer.eval()

    text_encoder = ContextCLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")

    unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
    vae.eval()
    text_encoder.eval()
    scheduler = PNDMScheduler(
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        set_alpha_to_one=False,
        skip_prk_steps=True,
    )
    tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
    image_processor = BlipImageProcessor()
    blip_diffusion = BlipDiffusionPipeline(
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        vae=vae,
        unet=unet,
        scheduler=scheduler,
        qformer=qformer,
        image_processor=image_processor,
    )
    blip_diffusion.save_pretrained(args.checkpoint_path)


def main(args):
    model, _, _ = load_model_and_preprocess("blip_diffusion", "base", device="cpu", is_eval=True)
    save_blip_diffusion_model(model.state_dict(), args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
    args = parser.parse_args()

    main(args)