import gradio as gr
import torch
from transformers import AutoModelForVisualQuestionAnswering

# Загружаем модель
model = AutoModelForVisualQuestionAnswering.from_pretrained("microsoft/visual-turing-nlg-question-answering")

# Создаем интерфейс gradio
def generate_answer(image, prompt=None):
    # Преобразуем изображение в тензор
    image = torch.as_tensor(image)

    # Если указан дополнительный prompt, добавляем его к запросу
    if prompt is not None:
        prompt = f"{prompt} {image.shape[1]} {image.shape[0]}"

    # Получаем ответ от модели
    answer = model(image, prompt=prompt).logits[0].argmax(dim=-1)

    # Возвращаем ответ
    return answer

# Создаем интерфейс
gr.Interface(generate_answer, inputs=[gr.Image(), gr.Text()], outputs=gr.Text())