Views
No views yet
1from __future__ import annotations
2
3import argparse
4import json
5from pathlib import Path
6from typing import Sequence
7
8import numpy as np
9import torch
10from PIL import Image
11from transformers import AutoModel, AutoProcessor
12
13
14STATE_DIM = 80
15IMAGE_SIZE = (448, 448)
16
17
18class MiniCPMVLAInference:
19 """Processor and model wrapper for single-sample VLA inference."""
20
21 def __init__(
22 self,
23 checkpoint_path: str | Path = "openbmb/MiniCPM-RobotManip",
24 device: str | torch.device | None = None,
25 ):
26 if device is None:
27 device = "cuda" if torch.cuda.is_available() else "cpu"
28 self.device = torch.device(device)
29 checkpoint = str(checkpoint_path)
30 self.processor = AutoProcessor.from_pretrained(checkpoint, trust_remote_code=True)
31 self.model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True)
32 self.model.to(self.device).eval()
33
34 @staticmethod
35 def _load_images(images: Sequence[str | Path | Image.Image | np.ndarray]) -> list[np.ndarray]:
36 if not images:
37 raise ValueError("At least one image is required")
38 loaded = []
39 for image in images:
40 if isinstance(image, (str, Path)):
41 with Image.open(image) as pil_image:
42 array = np.asarray(pil_image.convert("RGB"))
43 elif isinstance(image, Image.Image):
44 array = np.asarray(image.convert("RGB"))
45 elif isinstance(image, np.ndarray):
46 array = image
47 else:
48 raise TypeError(f"Unsupported image type: {type(image)!r}")
49 if array.ndim != 3 or array.shape[-1] != 3:
50 raise ValueError(f"Expected an HxWx3 image, got shape {array.shape}")
51 # Match the training pipeline's ResizeImage(size=(448, 448)),
52 # including PIL's default resize interpolation.
53 resized = Image.fromarray(array).resize(IMAGE_SIZE)
54 loaded.append(np.asarray(resized).copy())
55 return loaded
56
57 def preprocess(self, images: Sequence, text: str) -> dict[str, torch.Tensor]:
58 """Apply the same MiniCPM-V chat template and processor as training."""
59 content = [
60 {"type": "image", "image": image}
61 for image in self._load_images(images)
62 ]
63 content.append({"type": "text", "text": text})
64 messages = [{"role": "user", "content": content}]
65 inputs = self.processor.apply_chat_template(
66 messages,
67 tokenize=True,
68 add_generation_prompt=True,
69 return_dict=True,
70 return_tensors="pt",
71 processor_kwargs={"padding": False},
72 )
73 return {
74 key: value.to(self.device)
75 for key, value in inputs.items()
76 if isinstance(value, torch.Tensor)
77 }
78
79 def _prepare_state(self, state: torch.Tensor | np.ndarray | Sequence[float]) -> torch.Tensor:
80 state = torch.as_tensor(state, dtype=torch.float32, device=self.device)
81 if state.ndim == 1:
82 state = state.unsqueeze(0).unsqueeze(0)
83 elif state.ndim == 2:
84 state = state.unsqueeze(1)
85 if state.shape != (1, 1, STATE_DIM):
86 raise ValueError(f"state must have shape (80,), (1, 80), or (1, 1, 80); got {tuple(state.shape)}")
87 return state
88
89 @torch.inference_mode()
90 def predict(
91 self,
92 images: Sequence[str | Path | Image.Image | np.ndarray],
93 text: str,
94 state: torch.Tensor | np.ndarray | Sequence[float] | None = None,
95 embodiment_id: int = 0,
96 seed: int | None = None,
97 ) -> torch.Tensor:
98 """Return one action chunk with shape ``(30, 80)`` on CPU."""
99 if not 0 <= embodiment_id < self.model.action_head.max_num_embodiments:
100 raise ValueError(
101 f"embodiment_id must be in [0, {self.model.action_head.max_num_embodiments - 1}]"
102 )
103 if state is None:
104 state = torch.zeros(STATE_DIM)
105 state_tensor = self._prepare_state(state)
106 embodiment = torch.tensor([embodiment_id], dtype=torch.long, device=self.device)
107 if seed is not None:
108 torch.manual_seed(seed)
109 if self.device.type == "cuda":
110 torch.cuda.manual_seed_all(seed)
111
112 vlm_inputs = self.preprocess(images, text)
113 actions = self.model.predict_action(
114 state=state_tensor,
115 embodiment_id=embodiment,
116 **vlm_inputs,
117 )
118 return actions[0].float().cpu()
119
120
121def parse_args() -> argparse.Namespace:
122 parser = argparse.ArgumentParser(description=__doc__)
123 parser.add_argument("--image", action="append", required=True, help="Input image; repeat for multiple views")
124 parser.add_argument("--text", required=True, help="Robot instruction/prompt")
125 parser.add_argument("--device", default=None, help="Default: cuda if available, otherwise cpu")
126 state_group = parser.add_mutually_exclusive_group()
127 state_group.add_argument("--state-file", help="A .npy file containing 80 state values")
128 state_group.add_argument("--state", nargs=STATE_DIM, type=float, metavar="VALUE")
129 parser.add_argument("--embodiment-id", type=int, default=0)
130 parser.add_argument("--seed", type=int, default=None)
131 parser.add_argument("--output", help="Optional output .npy path; otherwise print JSON")
132 return parser.parse_args()
133
134
135if __name__ == "__main__":
136 args = parse_args()
137 if args.state_file:
138 state = np.load(args.state_file)
139 elif args.state is not None:
140 state = args.state
141 else:
142 state = np.zeros(STATE_DIM, dtype=np.float32)
143
144 infer_runner = MiniCPMVLAInference(
145 checkpoint_path="openbmb/MiniCPM-RobotManip",
146 device=args.device,
147 )
148 action = infer_runner.predict(
149 images=args.image,
150 text=args.text,
151 state=state,
152 embodiment_id=args.embodiment_id,
153 seed=args.seed,
154 )
155 if args.output:
156 output_path = Path(args.output)
157 output_path.parent.mkdir(parents=True, exist_ok=True)
158 np.save(output_path, action.numpy())
159 print(f"Saved action {tuple(action.shape)} to {output_path}")
160 else:
161 print(json.dumps(action.tolist()))
162