Views
No views yet
1from transformers import AutoProcessor
2import modeling_contextvla
3
4processor = AutoProcessor.from_pretrained("huiwon/ContextVLA-3B-Qwen2.5VL-FAST", use_fast=True)
5processor.tokenizer.padding_side = 'left'
6
7fast_tokenizer = AutoProcessor.from_pretrained(
8 "physical-intelligence/fast", trust_remote_code=True
9)
10# time_horizon and action_dim should be defined (task-specifically)
11fast_tokenizer.time_horizon = time_horizon
12fast_tokenizer.action_dim = action_dim
13
14model = modeling_contextvla.ContextVLA_Qwen2_5_VL.from_pretrained(
15 "huiwon/ContextVLA-3B-Qwen2.5VL-FAST",
16 attn_implementation="flash_attention_2",
17 dtype=torch.bfloat16
18)1import numpy as np
2
3def array_to_pil_image(self, frame):
4 if len(frame.shape) == 3 and frame.shape[0] == 3:
5 frame = np.transpose(frame, (1, 2, 0))
6 if frame.dtype != np.uint8:
7 if frame.max() <= 1.0:
8 frame = (frame * 255).astype(np.uint8)
9 else:
10 frame = frame.astype(np.uint8)
11 return Image.fromarray(frame)
12
13main_pixel_values = np.zeros((8, 224, 224, 3))
14wrist_pixel_values = np.zeros((8, 224, 224, 3))
15right_pixel_values = np.zeros((8, 224, 224, 3))
16
17pixel_values = np.stack([main_pixel_values, wrist_pixel_values, right_pixel_values], axis=1)
18pixel_values = pixel_values.reshape(-1, pixel_values.shape[-3], pixel_values.shape[-2], pixel_values.shape[-1])
19
20image_contents = [{"type": "image", "image": array_to_pil_image(frame)} for frame in pixel_values]
21
22messages = [
23 {
24 "role": "user",
25 "content": image_contents + [{"type": "text", "text": task_description}],
26 }
27]
28text = processor.apply_chat_template(
29 messages, tokenize=False, add_generation_prompt=True
30)
31
32image_inputs, video_inputs = vision_process.process_vision_info(messages)
33
34inputs = processor(
35 text=[text],
36 images=image_inputs,
37 videos=video_inputs,
38 padding=False,
39 return_tensors="pt",
40)1import torch
2
3model.model.layers[2].input_id_context = inputs['input_ids'].detach()
4ACTION_TOKEN_MIN = 151665
5ACTION_TOKEN_MAX = 153712
6
7# norm_stats q1 and q99 should be loaded
8action_high, action_low = np.array(norm_stats["norm_stats"]["actions"]["q99"]), np.array(norm_stats["norm_stats"]["actions"]["q01"])
9error_action = np.zeros((time_horizon, action_dim))
10
11with torch.no_grad():
12 generated_ids = model.generate(
13 **inputs,
14 max_new_tokens=256,
15 do_sample=False,
16 pad_token_id=processor.tokenizer.eos_token_id
17 )
18
19 action_indices = (ACTION_TOKEN_MIN <= generated_ids[0]) & (generated_ids[0] <= ACTION_TOKEN_MAX)
20 action_indices = torch.where(action_indices)[0]
21
22 output_action = fast_tokenizer.decode([generated_ids[0][action_indices] - ACTION_TOKEN_MIN])[0]
23
24 if np.allclose(error_action, output_action):
25 unnorm_actions = output_action
26 else:
27 unnorm_actions = (
28 0.5 * (output_action + 1) * (action_high - action_low)
29 + action_low
30 )
31 action = np.array(unnorm_actions)