MiniCPM-RobotManip

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor


STATE_DIM = 80
IMAGE_SIZE = (448, 448)


class MiniCPMVLAInference:
    """用于单样本VLA推理的处理器和模型封装器。"""

    def __init__(
        self,
        checkpoint_path: str | Path = "openbmb/MiniCPM-RobotManip",
        device: str | torch.device | None = None,
    ):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)
        checkpoint = str(checkpoint_path)
        self.processor = AutoProcessor.from_pretrained(checkpoint, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True)
        self.model.to(self.device).eval()

    @staticmethod
    def _load_images(images: Sequence[str | Path | Image.Image | np.ndarray]) -> list[np.ndarray]:
        if not images:
            raise ValueError("至少需要一张图像")
        loaded = []
        for image in images:
            if isinstance(image, (str, Path)):
                with Image.open(image) as pil_image:
                    array = np.asarray(pil_image.convert("RGB"))
            elif isinstance(image, Image.Image):
                array = np.asarray(image.convert("RGB"))
            elif isinstance(image, np.ndarray):
                array = image
            else:
                raise TypeError(f"不支持的图像类型: {type(image)!r}")
            if array.ndim != 3 or array.shape[-1] != 3:
                raise ValueError(f"期望HxWx3格式的图像，实际形状为 {array.shape}")
            # 匹配训练流程中的ResizeImage(size=(448, 448))，
            # 包括PIL默认的 resize 插值方式。
            resized = Image.fromarray(array).resize(IMAGE_SIZE)
            loaded.append(np.asarray(resized).copy())
        return loaded

    def preprocess(self, images: Sequence, text: str) -> dict[str, torch.Tensor]:
        """应用与训练时相同的MiniCPM-V对话模板和处理器。"""
        content = [
            {"type": "image", "image": image}
            for image in self._load_images(images)
        ]
        content.append({"type": "text", "text": text})
        messages = [{"role": "user", "content": content}]
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
            processor_kwargs={"padding": False},
        )
        return {
            key: value.to(self.device)
            for key, value in inputs.items()
            if isinstance(value, torch.Tensor)
        }

    def _prepare_state(self, state: torch.Tensor | np.ndarray | Sequence[float]) -> torch.Tensor:
        state = torch.as_tensor(state, dtype=torch.float32, device=self.device)
        if state.ndim == 1:
            state = state.unsqueeze(0).unsqueeze(0)
        elif state.ndim == 2:
            state = state.unsqueeze(1)
        if state.shape != (1, 1, STATE_DIM):
            raise ValueError(f"状态的形状必须为 (80,)、(1, 80) 或 (1, 1, 80)；实际为 {tuple(state.shape)}")
        return state

    @torch.inference_mode()
    def predict(
        self,
        images: Sequence[str | Path | Image.Image | np.ndarray],
        text: str,
        state: torch.Tensor | np.ndarray | Sequence[float] | None = None,
        embodiment_id: int = 0,
        seed: int | None = None,
    ) -> torch.Tensor:
        """返回一个形状为 ``(30, 80)`` 的动作序列，位于CPU上。"""
        if not 0 <= embodiment_id < self.model.action_head.max_num_embodiments:
            raise ValueError(
                f"embodiment_id 必须在 [0, {self.model.action_head.max_num_embodiments - 1}] 范围内"
            )
        if state is None:
            state = torch.zeros(STATE_DIM)
        state_tensor = self._prepare_state(state)
        embodiment = torch.tensor([embodiment_id], dtype=torch.long, device=self.device)
        if seed is not None:
            torch.manual_seed(seed)
            if self.device.type == "cuda":
                torch.cuda.manual_seed_all(seed)

        vlm_inputs = self.preprocess(images, text)
        actions = self.model.predict_action(
            state=state_tensor,
            embodiment_id=embodiment,
            **vlm_inputs,
        )
        return actions[0].float().cpu()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", action="append", required=True, help="输入图像；重复此参数可传入多视角图像")
    parser.add_argument("--text", required=True, help="机器人指令/提示词")
    parser.add_argument("--device", default=None, help="默认值：如果可用则为cuda，否则为cpu")
    state_group = parser.add_mutually_exclusive_group()
    state_group.add_argument("--state-file", help="包含80个状态值的.npy文件")
    state_group.add_argument("--state", nargs=STATE_DIM, type=float, metavar="VALUE")
    parser.add_argument("--embodiment-id", type=int, default=0)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--output", help="可选的输出.npy路径；若不指定则打印JSON")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if args.state_file:
        state = np.load(args.state_file)
    elif args.state is not None:
        state = args.state
    else:
        state = np.zeros(STATE_DIM, dtype=np.float32)

    infer_runner = MiniCPMVLAInference(
        checkpoint_path="openbmb/MiniCPM-RobotManip",
        device=args.device,
    )
    action = infer_runner.predict(
        images=args.image,
        text=args.text,
        state=state,
        embodiment_id=args.embodiment_id,
        seed=args.seed,
    )
    if args.output:
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(output_path, action.numpy())
        print(f"已将形状为 {tuple(action.shape)} 的动作保存至 {output_path}")
    else:
        print(json.dumps(action.tolist()))