import torch
import gradio as gr
from promptcap import PromptCap

# モデルの読み込み
model = PromptCap("tifa-benchmark/promptcap-coco-vqa")

# CUDAが使用可能ならGPUを使用
if torch.cuda.is_available():
    model.cuda()

# 画像と質問を入力として、キャプションを生成する関数を定義
def generate_caption(image, question):
    prompt = f"please describe this image according to the given question: {question}"
    # PromptCapモデルでキャプションを生成
    caption = model.caption(prompt, image)
    return caption

# Gradioインターフェースの定義
interface = gr.Interface(
    fn=generate_caption,  # キャプション生成関数
    inputs=[
        gr.inputs.Image(type="filepath", label="Input Image"),  # 画像入力
        gr.inputs.Textbox(label="Question")  # 質問入力
    ],
    outputs=gr.outputs.Textbox(label="Generated Caption")  # キャプション出力
)

# インターフェースを起動
interface.launch()