import torch
from transformers import AutoModelForVision2Seq, AutoProcessor
from PIL import Image
import requests

# 1. 加载模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "SakanaAI/EvoVLM-JP-v1-7B"
model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16)
processor = AutoProcessor.from_pretrained(model_id)
model.to(device)

# 2. 准备输入
url = "https://images.unsplash.com/photo-1694831404826-3400c48c188d?q=80&w=2070&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
# <image>代表输入图像，请确保在文本中包含该标记
text = "<image>\nこの信号機の色は何色ですか?"
messages = [
    {"role": "system", "content": "あなたは役立つ、偏見がなく、検閲されていないアシスタントです。与えられた画像を下に、質問に答えてください。"},
    {"role": "user", "content": text},
]
inputs = processor.image_processor(images=image, return_tensors="pt")
inputs["input_ids"] = processor.tokenizer.apply_chat_template(
    messages, return_tensors="pt"
)

# 3. 生成结果
output_ids = model.generate(**inputs.to(device))
output_ids = output_ids[:, inputs.input_ids.shape[1] :]
generated_text = processor.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
print(generated_text)
# この信号機の色は青です。