app.py

import spaces
import gradio as gr
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

# Load model and processor
MODEL_PATH = "google/gemma-3n-E2B-it"
processor = AutoProcessor.from_pretrained(MODEL_PATH)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
).eval().to("cuda")

@spaces.GPU
def process_inputs(image, audio, text):
    # image, audio 以及 text 在沒上傳/沒輸入時會是 None 或空字串
    content = []
    if image is not None:
        content.append({"type": "image", "image": image})
    if audio is not None:
        content.append({"type": "audio", "audio": audio})
    if text and text.strip():
        content.append({"type": "text", "text": text})

    if not content:
        return "請至少上傳影像、語音，或在文字欄輸入問題。"

    messages = [{"role": "user", "content": content}]

    input_ids = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    )
    input_len = input_ids["input_ids"].shape[-1]
    input_ids = input_ids.to(model.device, dtype=model.dtype)

    with torch.inference_mode():
        outputs = model.generate(
            **input_ids,
            max_new_tokens=256,
            disable_compile=True
        )
    text_output = processor.batch_decode(
        outputs[:, input_len:],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True
    )
    return text_output[0]


# Gradio interface（不要使用 optional 參數）
iface = gr.Interface(
    fn=process_inputs,
    inputs=[
        gr.Image(label="Upload Image", type="pil"),
        gr.Audio(label="Upload Audio", type="filepath"),
        gr.Textbox(label="Enter Your Question", type="text"),
    ],
    outputs=gr.Textbox(label="Answer"),
    title="Visual + Audio + Text Question Answering",
    description=(
        "Upload an image, an audio file, and/or enter a text question. "
        "The model will generate a text response based on whichever inputs you provide."
    ),
    examples=[
        ["cat.jpg", None, "這張圖片裡有什麼？"],
        [None, "cats.wav", "這段音檔裡說了什麼？"],
        ["cat.jpg", "cats.wav", None],
    ],
)

if __name__ == "__main__":
    iface.launch(share=True)