Views
No views yet
transformers >= 4.47.0, you can directly use the following code to load the model and perform inference.1
2from typing import Dict, List, Optional, Union
3from pathlib import Path
4import numpy as np
5import torch
6from PIL import Image
7from transformers import LlamaTokenizerFast
8from transformers import AutoConfig, AutoImageProcessor, AutoModelForVision2Seq, AutoProcessor
9import os, json
10from model.prismatic import PrismaticVLM
11from model.overwatch import initialize_overwatch
12from model.action_tokenizer import ActionTokenizer
13from model.vision_backbone import DinoSigLIPViTBackbone, DinoSigLIPImageTransform
14from model.llm_backbone import LLaMa2LLMBackbone
15from extern.hf.configuration_prismatic import OpenFlyConfig
16from extern.hf.modeling_prismatic import OpenVLAForActionPrediction
17from extern.hf.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor
18
19AutoConfig.register("openvla", OpenFlyConfig)
20AutoImageProcessor.register(OpenFlyConfig, PrismaticImageProcessor)
21AutoProcessor.register(OpenFlyConfig, PrismaticProcessor)
22AutoModelForVision2Seq.register(OpenFlyConfig, OpenVLAForActionPrediction)
23
24model_name_or_path="IPEC-COMMUNITY/openfly-agent-7b"
25processor = AutoProcessor.from_pretrained(model_name_or_path)
26model = AutoModelForVision2Seq.from_pretrained(
27 model_name_or_path,
28 attn_implementation="flash_attention_2", # [Optional] Requires `flash_attn`
29 torch_dtype=torch.bfloat16,
30 low_cpu_mem_usage=True,
31 trust_remote_code=True,
32).to("cuda:0")
33
34image = Image.fromarray(cv2.imread("example.png"))
35prompt = "Take off, go straight pass the river"
36inputs = processor(prompt, [image, image, image]).to("cuda:0", dtype=torch.bfloat16)
37action = model.predict_action(**inputs, unnorm_key="vln_norm", do_sample=False)
38print(action)