import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor
model = AutoModelForImageTextToText.from_pretrained(
"AIcell/guava-05-22", torch_dtype=torch.bfloat16, device_map="cuda",
)
proc = AutoProcessor.from_pretrained("AIcell/guava-05-22")
system_prompt = open("system_prompt.txt").read().strip()
scene_img = Image.open("scene.png").convert("RGB")
messages = [
{"role": "system", "content": [{"type": "text", "text": system_prompt}]},
{"role": "user", "content": [
{"type": "image", "image": scene_img},
{"type": "text", "text":
"Task: <your task description>.\n\n"
"Gripper is at [...] rotation [...] width X%."},
]},
]
inputs = proc.apply_chat_template(
messages, add_generation_prompt=True,
tokenize=True, return_dict=True, return_tensors="pt",
).to("cuda")
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=512, do_sample=False)
print(proc.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0])