import requests
import torch
from transformers import AutoModelForCausalLM
from PIL import Image
import gradio as gr

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    "q-future/one-align", 
    trust_remote_code=True, 
    attn_implementation="eager", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

def score_image(image, task):
    """
    对输入的图像进行评分
    :param image: 输入的图像
    :param task: 任务类型，可以是 "quality" 或 "aesthetics"
    :return: 评分结果
    """
    if task not in ["quality", "aesthetics"]:
        return "任务类型必须是 'quality' 或 'aesthetics'"
    
    # 将图像转换为模型所需的格式
    if isinstance(image, str):
        image = Image.open(requests.get(image, stream=True).raw)
    elif isinstance(image, Image.Image):
        pass
    else:
        return "输入必须是图像 URL 或 PIL 图像"
    
    # 调用模型进行评分
    result = model.score([image], task_=task, input_="image")
    return result

# 创建 Gradio 界面
iface = gr.Interface(
    fn=score_image, 
    inputs=[
        gr.Image(label="输入图像", type="pil"), 
        gr.Dropdown(choices=["quality", "aesthetics"], label="任务类型")
    ], 
    outputs="text",
    title="图像评分模型",
    description="上传图像并选择任务类型（quality 或 aesthetics）来获取评分。"
)

# 启动 Gradio 应用
iface.launch()