import gradio as gr 
import gradio.components as grc
import torch

from lavis.models import load_model_and_preprocess
from lavis.processors import load_processor

# setup device to use
device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
model, vis_processors, text_processors = load_model_and_preprocess("blip2_image_text_matching", "pretrain", device=device, is_eval=True)

def predict(raw_image, caption):
    raw_image = raw_image.convert("RGB")
    img = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
    txt = text_processors["eval"](caption)
    itm_output = model({"image": img, "text_input": txt}, match_head="itm")
    itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
    itm_score = itm_scores[:, 1].item()
    itc_score = model({"image": img, "text_input": txt}, match_head='itc')
    return '%.3f' % itm_score, '%.4f' % itc_score

app = gr.Interface(fn=predict, inputs=[grc.Image(type="pil"), grc.Textbox()], outputs=[grc.Text(label="itm score"), grc.Text(label="itc score")])
app.launch()