import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
from vqvae import MultiVQVAE
from models.normalizer import LinearNormalizer
from utils import batch_predict_action
device = "cuda:0"
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
"robotics-diffusion-transformer/RDT2-VQ"
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
device_map=device
).eval()
vae = MultiVQVAE.from_pretrained("robotics-diffusion-transformer/RVQActionTokenizer").eval()
vae = vae.to(device=device, dtype=torch.float32)
valid_action_id_length = (
vae.pos_id_len + vae.rot_id_len + vae.grip_id_len
)
normalizer = LinearNormalizer.from_pretrained("umi_normalizer_wo_downsample_indentity_rot.pt")
result = batch_predict_action(
model,
processor,
vae,
normalizer,
examples=[
{
"obs": {
"camera0_rgb": ...,
"camera1_rgb": ...,
},
"meta": {
"num_camera": 2
}
},
...,
],
valid_action_id_length=valid_action_id_length,
apply_jpeg_compression=True,
instruction="Pick up the apple."
)
action_chunk = result["action_pred"][0]
for robot_idx in range(2):
action_chunk[:, robot_idx * 10 + 9] = action_chunk[:, robot_idx * 10 + 9] / 0.088 * 0.1