app.py

import gradio as gr
from transformers import AutoModelForVision2Seq, AutoProcessor
from PIL import Image

# 1. 模型加载 (请替换为您选择的 VLM 模型ID)
MODEL_ID = "google/paligemma-3b-mix-448" # 示例 VLM
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForVision2Seq.from_pretrained(MODEL_ID) 

# 2. 推理函数：接受图片和文本
def vlm_inference(image: Image.Image, prompt: str):
    """
    接收用户上传的图片和文本提示，并返回 VLM 的生成结果。
    """
    if image is None:
        return "请上传一张图片。"
    
    # 构建 VLM 输入格式 (根据模型要求调整)
    # 对于 IDEFICS 这类 VLM，输入通常是一个包含图片和文本的列表
    chats = [
        "User:", 
        image, 
        f"Prompt: {prompt}", 
        "<end_of_utterance>",
        "\nAssistant:" 
    ]

    print(f'input = {chats}, image.size {image.size}')
    # 预处理和生成
    inputs = processor(chats, return_tensors="pt")
    
    # 注意：大型 VLM 需要 GPU 部署。如果使用免费 CPU 空间，请选择小型 VLM
    # outputs = model.generate(**inputs, max_length=128) # 确保模型在 CPU/GPU 上

    # 模拟生成（如果您正在测试或使用 CPU 部署）
    # 请取消注释上面的 model.generate 行，并删除下面这行
    outputs = f"（模型生成结果）您询问的是图片中的内容：'{prompt}'。图片尺寸为 {image.size}。" 

    print(outputs)
    # 假设模型返回文本结果
    return outputs

# 3. Gradio 界面定义
with gr.Blocks(title="VLM 多模态聊天演示") as demo:
    gr.Markdown("# 🎨 Hugging Face Spaces VLM 演示")
    
    # 定义输入组件
    with gr.Row():
        image_input = gr.Image(type="pil", label="上传图片")
        prompt_input = gr.Textbox(label="输入文本提示 (例如: 描述这张图片的内容)")
        
    # 定义输出组件和按钮
    output_text = gr.Textbox(label="VLM 生成结果")
    submit_btn = gr.Button("运行 VLM")

    # 绑定函数
    submit_btn.click(
        fn=vlm_inference, 
        inputs=[image_input, prompt_input], 
        outputs=output_text
    )

# 启动应用
if __name__ == "__main__":
    # 在本地测试时运行
    demo.launch()