lampe-sim-data-openvla

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

# Load model and processor
processor = AutoProcessor.from_pretrained("kavinrajkrupsurge/lampe-sim-data-openvla", trust_remote_code=True)
model = AutoModelForVision2Seq.from_pretrained(
    "kavinrajkrupsurge/lampe-sim-data-openvla",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
).to("cuda")

# Prepare inputs
instruction = "search for a person in the room by scanning the room and stop when you find"
image = Image.open("path/to/image.jpg")
prompt = f"In: What action should the robot take to {instruction.lower()}?\nOut:"

inputs = processor(prompt, image).to("cuda", dtype=torch.bfloat16)

# Predict action
with torch.inference_mode():
    action = model.predict_action(
        **inputs,
        unnorm_key="lampe_search_dataset/all",
        do_sample=False
    )