import torch
from transformers import AutoTokenizer, AutoModelForVision2Seq, AutoImageProcessor
from PIL import Image
import requests

# 构建输入提示的辅助函数
TASK2INSTRUCTION = {
    "caption": "请详细描述这张图片。",
    "tag": "使用给定的词语详细描述这张图片。",
    "vqa": "根据提供的图片回答问题。",
}


def build_prompt(task="caption", input=None, sep="\n\n### "):
    assert (
        task in TASK2INSTRUCTION
    ), f"请从{list(TASK2INSTRUCTION.keys())}中选择任务"
    if task in ["tag", "vqa"]:
        assert input is not None, "请填写`input`参数！"
        if task == "tag" and isinstance(input, list):
            input = "、".join(input)
    else:
        assert input is None, f"`{task}`模式不支持输入问题"
    sys_msg = "以下是指令说明与上下文输入组合的任务。请撰写能恰当满足要求的回答。"
    p = sys_msg
    roles = ["指令", "回答"]
    instruction = TASK2INSTRUCTION[task]
    msgs = [": \n" + instruction, ": \n"]
    if input:
        roles.insert(1, "输入")
        msgs.insert(1, ": \n" + input)
    for role, msg in zip(roles, msgs):
        p += sep + role + msg
    return p

# 加载模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForVision2Seq.from_pretrained("stabilityai/japanese-stable-vlm", trust_remote_code=True)
processor = AutoImageProcessor.from_pretrained("stabilityai/japanese-stable-vlm")
tokenizer = AutoTokenizer.from_pretrained("stabilityai/japanese-stable-vlm")
model.to(device)

# 准备输入
url = "https://images.unsplash.com/photo-1582538885592-e70a5d7ab3d3?ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D&auto=format&fit=crop&w=1770&q=80"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
prompt = build_prompt(task="caption")
# prompt = build_prompt(task="tag", input=["河津樱", "蓝天"])
# prompt = build_prompt(task="vqa", input="这是什么季节？")

inputs = processor(images=image, return_tensors="pt")
text_encoding = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
inputs.update(text_encoding)

# 生成结果
outputs = model.generate(
    **inputs.to(device, dtype=model.dtype),
    do_sample=False,
    num_beams=5,
    max_new_tokens=128,
    min_length=1,
    repetition_penalty=1.5,
)
generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].strip()
print(generated_text)
# 樱花映衬下的东京晴空塔