import gradio as gr
from transformers import ViltProcessor, ViltForQuestionAnswering
import torch

torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
torch.hub.download_url_to_file('https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg', 'astronaut.jpg')

processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")

def getAnswer(image,text):
    encoding = processor(image, text, return_tensors="pt")
    
    # forward pass
    with torch.no_grad():
     outputs = model(**encoding)
     
    logits = outputs.logits
    idx = logits.argmax(-1).item()
    predicted_answer = model.config.id2label[idx]
   
    return predicted_answer

image = gr.inputs.Image(type="pil")
question = gr.inputs.Textbox(label="Question about the image")
answer = gr.outputs.Textbox(label="Predicted answer")
examples = [["cats.jpg", "How many cats are there?"], ["astronaut.jpg", "What's the astronaut riding on?"]]


title="Visual question and answering"

iface = gr.Interface(fn=getAnswer, 
                         inputs=[image, question], 
                         outputs=answer, 
                         examples=examples, 
                         title=title,
                         enable_queue=True)
iface.launch(debug=True )