Student-gpt / app.py
Rooni's picture
Update app.py
ba5e8fe
raw
history blame contribute delete
921 Bytes
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())