import gradio as gr
import torch
from transformers import (AutoProcessor, BlipForQuestionAnswering,
                          ViltForQuestionAnswering)

torch.hub.download_url_to_file(
    'http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
torch.hub.download_url_to_file(
    'https://huggingface.co/datasets/nielsr/textcaps-sample/resolve/main/stop_sign.png',
    'stop_sign.png')
torch.hub.download_url_to_file(
    'https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg',
    'astronaut.jpg')

blip_processor_large = AutoProcessor.from_pretrained(
    'Salesforce/blip-vqa-capfilt-large')
blip_model_large = BlipForQuestionAnswering.from_pretrained(
    'Salesforce/blip-vqa-capfilt-large')

vilt_processor = AutoProcessor.from_pretrained(
    'dandelin/vilt-b32-finetuned-vqa')
vilt_model = ViltForQuestionAnswering.from_pretrained(
    'dandelin/vilt-b32-finetuned-vqa')

device = 'cuda' if torch.cuda.is_available() else 'cpu'

blip_model_large.to(device)
vilt_model.to(device)


@torch.inference_mode()
def generate_answer_blip(processor, model, image, question):
    inputs = processor(images=image, text=question,
                       return_tensors='pt').to(device)
    generated_ids = model.generate(**inputs, max_length=50)
    generated_answer = processor.batch_decode(generated_ids,
                                              skip_special_tokens=True)
    return generated_answer[0]


@torch.inference_mode()
def generate_answer_vilt(processor, model, image, question):
    encoding = processor(images=image, text=question,
                         return_tensors='pt').to(device)
    outputs = model(**encoding)
    predicted_class_idx = outputs.logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]


def generate_answers(image, question):
    answer_blip_large = generate_answer_blip(blip_processor_large,
                                             blip_model_large, image, question)
    answer_vilt = generate_answer_vilt(vilt_processor, vilt_model, image,
                                       question)
    return answer_blip_large, answer_vilt


demo = gr.Interface(
    fn=generate_answers,
    inputs=[gr.Image(type='pil'),
            gr.Textbox(label='Question')],
    outputs=[
        gr.Textbox(label='Answer generated by BLIP-large'),
        gr.Textbox(label='Answer generated by ViLT')
    ],
    examples=[
        ['cats.jpg', 'How many cats are there?'],
        ['stop_sign.png', "What's behind the stop sign?"],
        ['astronaut.jpg', "What's the astronaut riding on?"],
    ],
    title='Interactive demo: comparing visual question answering (VQA) models')
demo.queue().launch()