Visual Question Answering and Image Captioning using BLIP and OpenVINO#

import torch
from pathlib import Path
import openvino as ov

VISION_MODEL_OV = Path("blip_vision_model.xml")
vision_model = model.vision_model
vision_model.eval()

# check that model works and save it outputs for reusage as text encoder input
with torch.no_grad():
    vision_outputs = vision_model(inputs["pixel_values"])

# if openvino model does not exist, convert it to IR
if not VISION_MODEL_OV.exists():
    # export pytorch model to ov.Model
    with torch.no_grad():
        ov_vision_model = ov.convert_model(vision_model, example_input=inputs["pixel_values"])
    # save model on disk for next usages
    ov.save_model(ov_vision_model, VISION_MODEL_OV)
    print(f"Vision model successfuly converted and saved to {VISION_MODEL_OV}")
else:
    print(f"Vision model will be loaded from {VISION_MODEL_OV}")

TEXT_ENCODER_OV = Path("blip_text_encoder.xml")


text_encoder = model.text_encoder
text_encoder.eval()

# if openvino model does not exist, convert it to IR
if not TEXT_ENCODER_OV.exists():
    # prepare example inputs
    image_embeds = vision_outputs[0]
    image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long)
    input_dict = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "encoder_hidden_states": image_embeds,
        "encoder_attention_mask": image_attention_mask,
    }
    # export PyTorch model
    with torch.no_grad():
        ov_text_encoder = ov.convert_model(text_encoder, example_input=input_dict)
    # save model on disk for next usages
    ov.save_model(ov_text_encoder, TEXT_ENCODER_OV)
    print(f"Text encoder successfuly converted and saved to {TEXT_ENCODER_OV}")
else:
    print(f"Text encoder will be loaded from {TEXT_ENCODER_OV}")

text_decoder = model.text_decoder
text_decoder.eval()

TEXT_DECODER_OV = Path("blip_text_decoder_with_past.xml")

# prepare example inputs
input_ids = torch.tensor([[30522]])  # begin of sequence token id
attention_mask = torch.tensor([[1]])  # attention mask for input_ids
encoder_hidden_states = torch.rand((1, 10, 768))  # encoder last hidden state from text_encoder
encoder_attention_mask = torch.ones((1, 10), dtype=torch.long)  # attention mask for encoder hidden states

input_dict = {
    "input_ids": input_ids,
    "attention_mask": attention_mask,
    "encoder_hidden_states": encoder_hidden_states,
    "encoder_attention_mask": encoder_attention_mask,
}
text_decoder_outs = text_decoder(**input_dict)
# extend input dictionary with hidden states from previous step
input_dict["past_key_values"] = text_decoder_outs["past_key_values"]

text_decoder.config.torchscript = True
if not TEXT_DECODER_OV.exists():
    # export PyTorch model
    with torch.no_grad():
        ov_text_decoder = ov.convert_model(text_decoder, example_input=input_dict)
    # save model on disk for next usages
    ov.save_model(ov_text_decoder, TEXT_DECODER_OV)
    print(f"Text decoder successfuly converted and saved to {TEXT_DECODER_OV}")
else:
    print(f"Text decoder will be loaded from {TEXT_DECODER_OV}")

# create OpenVINO Core object instance
core = ov.Core()

device = device_widget()

device

# load models on device
ov_vision_model = core.compile_model(VISION_MODEL_OV, device.value)
ov_text_encoder = core.compile_model(TEXT_ENCODER_OV, device.value)
ov_text_decoder_with_past = core.compile_model(TEXT_DECODER_OV, device.value)

from functools import partial
from blip_model import text_decoder_forward

text_decoder.forward = partial(text_decoder_forward, ov_text_decoder_with_past=ov_text_decoder_with_past)

if not Path("./blip_model.py").exists():
    download_file(url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/blip-visual-language-processing/blip_model.py")
from blip_model import OVBlipModel

ov_model = OVBlipModel(model.config, model.decoder_start_token_id, ov_vision_model, ov_text_encoder, text_decoder)
out = ov_model.generate_answer(**inputs, max_length=20)

out = ov_model.generate_caption(inputs["pixel_values"], max_length=20)
caption = processor.decode(out[0], skip_special_tokens=True)
fig = visualize_results(raw_image, caption)

start = time.perf_counter()
out = ov_model.generate_answer(**inputs, max_length=20)
end = time.perf_counter() - start
answer = processor.decode(out[0], skip_special_tokens=True)
fig = visualize_results(raw_image, answer, question)

print(f"Processing time: {end:.4f}")

%%skip not $to_quantize.value

import numpy as np
from datasets import load_dataset
from tqdm.notebook import tqdm

def preprocess_batch(batch, vision_model, inputs_info):
    """
    Preprocesses a dataset batch by loading and transforming image and text data.
    VQAv2 dataset contains multiple questions to image.
    To reduce dataset preparation time we will store preprocessed images in `inputs_info`.
    """
    image_id = batch["image_id"]
    if image_id in inputs_info:
        inputs = processor(text=batch['question'], return_tensors="np")
        pixel_values = inputs_info[image_id]["pixel_values"]
        encoder_hidden_states = inputs_info[image_id]["encoder_hidden_states"]
    else:
        inputs = processor(images=batch["image"], text=batch["question"], return_tensors="np")
        pixel_values = inputs["pixel_values"]
        encoder_hidden_states = vision_model(pixel_values)[vision_model.output(0)]
        inputs_info[image_id] = {
            "pixel_values": pixel_values,
            "encoder_hidden_states": encoder_hidden_states,
            "text_encoder_inputs": []
        }

    text_encoder_inputs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"]
    }
    inputs_info[image_id]["text_encoder_inputs"].append(text_encoder_inputs)


def prepare_input_data(dataloader, vision_model, opt_init_steps):
    """
    Store calibration subset in List to reduce quantization time.
    """
    inputs_info = {}
    for idx, batch in enumerate(tqdm(dataloader, total=opt_init_steps, desc="Prepare calibration data")):
        preprocess_batch(batch, vision_model, inputs_info)

    calibration_subset = []
    for image_id in inputs_info:
        pixel_values = inputs_info[image_id]["pixel_values"]
        encoder_hidden_states = inputs_info[image_id]["encoder_hidden_states"]
        encoder_attention_mask = np.ones(encoder_hidden_states.shape[:-1], dtype=int)
        for text_encoder_inputs in inputs_info[image_id]["text_encoder_inputs"]:
            text_encoder_inputs["encoder_hidden_states"] = encoder_hidden_states
            text_encoder_inputs["encoder_attention_mask"] = encoder_attention_mask
            blip_inputs = {
                "vision_model_inputs": {"pixel_values": pixel_values},
                "text_encoder_inputs": text_encoder_inputs,
            }
            calibration_subset.append(blip_inputs)
    return calibration_subset


def prepare_dataset(vision_model, opt_init_steps=300, streaming=False):
    """
    Prepares a vision-text dataset for quantization.
    """
    split = f"train[:{opt_init_steps}]" if not streaming else "train"
    dataset = load_dataset("HuggingFaceM4/VQAv2", split=split, streaming=streaming, trust_remote_code=True)
    dataset = dataset.shuffle(seed=42)
    if streaming:
        dataset = dataset.take(opt_init_steps)
    calibration_subset = prepare_input_data(dataset, vision_model, opt_init_steps)
    return calibration_subset

%%skip not $to_quantize.value

import nncf

comp_vision_model = core.compile_model(VISION_MODEL_OV, device.value)
calibration_data = prepare_dataset(comp_vision_model)

%%skip not $to_quantize.value

vision_dataset = nncf.Dataset(calibration_data, lambda x: x["vision_model_inputs"])
vision_model = core.read_model(VISION_MODEL_OV)

quantized_model = nncf.quantize(
    model=vision_model,
    calibration_dataset=vision_dataset,
    model_type=nncf.ModelType.TRANSFORMER
)

ov.save_model(quantized_model, VISION_MODEL_OV_INT8)

%%skip not $to_quantize.value

text_encoder_dataset = nncf.Dataset(calibration_data, lambda x: x["text_encoder_inputs"])
text_encoder_model = core.read_model(TEXT_ENCODER_OV)

quantized_model = nncf.quantize(
    model=text_encoder_model,
    calibration_dataset=text_encoder_dataset,
    model_type=nncf.ModelType.TRANSFORMER
)
ov.save_model(quantized_model, TEXT_ENCODER_OV_INT8)

%%skip not $to_quantize.value

text_decoder_model = core.read_model(TEXT_DECODER_OV)
compressed_text_decoder = nncf.compress_weights(text_decoder_model)
ov.save_model(compressed_text_decoder, str(TEXT_DECODER_OV_INT8))

%%skip not $to_quantize.value

q_ov_vision_model = core.compile_model(VISION_MODEL_OV_INT8, device.value)
q_ov_text_encoder = core.compile_model(TEXT_ENCODER_OV_INT8, device.value)
q_ov_text_decoder_with_past = core.compile_model(TEXT_DECODER_OV_INT8, device.value)

%%skip not $to_quantize.value

from functools import partial
from transformers import BlipForQuestionAnswering
from blip_model import OVBlipModel, text_decoder_forward

model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
text_decoder = model.text_decoder
text_decoder.eval()

text_decoder.forward = partial(text_decoder_forward, ov_text_decoder_with_past=q_ov_text_decoder_with_past)
int8_model = OVBlipModel(model.config, model.decoder_start_token_id, q_ov_vision_model, q_ov_text_encoder, text_decoder)

%%skip not $to_quantize.value

raw_image = Image.open("demo.jpg").convert('RGB')
question = "how many dogs are in the picture?"
# preprocess input data
inputs = processor(raw_image, question, return_tensors="pt")

%%skip not $to_quantize.value

def calculate_compression_rate(ov_model_path):
    fp16_ir_model_size = Path(ov_model_path).with_suffix(".bin").stat().st_size / 1024
    int8_model_path = str(ov_model_path).replace(".xml", "_int8.xml")
    quantized_model_size = Path(int8_model_path).with_suffix(".bin").stat().st_size / 1024
    print(f'{ov_model_path.as_posix().split(".")[0]}')
    print(f"    * FP16 IR model size: {fp16_ir_model_size:.2f} KB")
    print(f"    * INT8 model size: {quantized_model_size:.2f} KB")
    print(f"    * Model compression rate: {fp16_ir_model_size / quantized_model_size:.3f}")

%%skip not $to_quantize.value

for fp16_path in [VISION_MODEL_OV, TEXT_ENCODER_OV, TEXT_DECODER_OV]:
    calculate_compression_rate(fp16_path)

%%skip not $to_quantize.value

import time
import torch

def calculate_inference_time(blip_model, calibration_data, generate_caption):
    inference_time = []
    for inputs in calibration_data:
        pixel_values = torch.from_numpy(inputs["vision_model_inputs"]["pixel_values"])
        input_ids = torch.from_numpy(inputs["text_encoder_inputs"]["input_ids"])
        attention_mask = torch.from_numpy(inputs["text_encoder_inputs"]["attention_mask"])

        start = time.perf_counter()
        if generate_caption:
            _ = blip_model.generate_caption(pixel_values, max_length=20)
        else:
            _ = blip_model.generate_answer(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask, max_length=20)
        end = time.perf_counter()
        delta = end - start
        inference_time.append(delta)
    return np.median(inference_time)

%%skip not $to_quantize.value

fp_original_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
fp_text_decoder = fp_original_model.text_decoder
fp_text_decoder.eval()

comp_text_encoder = core.compile_model(TEXT_ENCODER_OV, device.value)
comp_text_decoder_with_past = core.compile_model(TEXT_DECODER_OV, device.value)
fp_text_decoder.forward = partial(text_decoder_forward, ov_text_decoder_with_past=comp_text_decoder_with_past)
fp16_model = OVBlipModel(model.config, model.decoder_start_token_id, comp_vision_model, comp_text_encoder, fp_text_decoder)

%%skip not $to_quantize.value

validation_data = calibration_data[:100]

int8_caption_latency = calculate_inference_time(int8_model, validation_data, generate_caption=True)
fp16_caption_latency = calculate_inference_time(fp16_model, validation_data, generate_caption=True)

print(f"Image Captioning speed up: {fp16_caption_latency / int8_caption_latency:.3f}")

%%skip not $to_quantize.value

int8_generate_answer_latency = calculate_inference_time(int8_model, validation_data, generate_caption=False)
fp16_generate_answer_latency = calculate_inference_time(fp16_model, validation_data, generate_caption=False)
print(f"Question Answering speed up: {fp16_generate_answer_latency / int8_generate_answer_latency:.3f}")