248 lines
11 KiB
Python
248 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""π0 (Pi-Zero) VLA pipeline for vllm-omni.
|
|
|
|
Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the
|
|
DreamZero contract: the pipeline owns ALL preprocessing. It reads the raw robot
|
|
observation from ``req.sampling_params.extra_args["robot_obs"]`` (delivered by
|
|
the OpenPI realtime serving layer), builds model inputs, runs flow-matching
|
|
denoising, and returns ``DiffusionOutput(output={"actions": ndarray})`` — which
|
|
``diffusion_engine`` promotes to ``multimodal_output["actions"]`` for the client.
|
|
|
|
π0 is stateless across calls (no KV reuse), so ``session_id`` / ``reset`` from
|
|
the OpenPI protocol are accepted but ignored.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.models.pi0.config import Pi0Config
|
|
from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction
|
|
from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs
|
|
from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE as PI0_PIPELINE
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Default tokenizer for the PaliGemma prefix (matches LeRobot Pi0).
|
|
DEFAULT_PI0_TOKENIZER = "google/paligemma-3b-pt-224"
|
|
|
|
|
|
def _pi0_post_process(x):
|
|
"""Module-level identity post-process (picklable across the orchestrator's
|
|
multiprocess boundary — a local closure is not)."""
|
|
return x
|
|
|
|
|
|
def get_pi0_post_process_func(od_config: OmniDiffusionConfig):
|
|
"""π0 returns actions directly; post-processing is identity. The diffusion
|
|
engine resolves this via the registry (``_DIFFUSION_POST_PROCESS_FUNCS``) and
|
|
applies it engine-side, so the pipeline does NOT attach it to DiffusionOutput.
|
|
"""
|
|
del od_config
|
|
return _pi0_post_process
|
|
|
|
|
|
class Pi0Pipeline(nn.Module):
|
|
"""π0 VLA pipeline: raw robot obs → continuous action chunk.
|
|
|
|
Registered as ``"Pi0Pipeline"`` in the diffusion registry. Weights are
|
|
self-loaded in ``__init__`` from the checkpoint's ``model.safetensors`` via
|
|
the kernel's ``load_weights`` (which handles the LeRobot key remaps).
|
|
"""
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
self.prefix = prefix
|
|
# Resolve od_config.model to a LOCAL directory. A bare HF repo id (e.g.
|
|
# the documented ``lerobot/pi0_base``) must be snapshot-downloaded first;
|
|
# otherwise config/tokenizer/weights would silently fall back to random
|
|
# init (this pipeline self-loads from a local dir and has no
|
|
# ``weights_sources`` for the framework loader to download from).
|
|
self.model_dir = self._resolve_model_dir(od_config.model)
|
|
self.config = self._build_config(od_config)
|
|
|
|
custom_args = od_config.custom_pipeline_args or {}
|
|
self.tokenizer_source = str(custom_args.get("tokenizer", self._resolve_tokenizer_source()))
|
|
|
|
# Torch dtype/device from od_config.
|
|
self._torch_dtype = self._resolve_dtype(od_config)
|
|
self._device = self._resolve_device(od_config)
|
|
|
|
self.tokenizer = self._load_tokenizer()
|
|
self.model = self._initialize_model()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Construction helpers
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _resolve_model_dir(model: str | None) -> str | None:
|
|
"""Return a local directory for ``model``; download an HF repo id if needed.
|
|
|
|
Limits the snapshot to the files π0 actually reads (config, weights,
|
|
tokenizer) so a repo id doesn't pull unrelated artifacts.
|
|
"""
|
|
if not model:
|
|
return None
|
|
if os.path.isdir(model):
|
|
return model
|
|
from huggingface_hub import snapshot_download
|
|
|
|
return snapshot_download(
|
|
repo_id=model,
|
|
allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"],
|
|
)
|
|
|
|
def _build_config(self, od_config: OmniDiffusionConfig) -> Pi0Config:
|
|
"""Build Pi0Config from deploy-yaml model_config, falling back to the
|
|
checkpoint's config.json (raw LeRobot format)."""
|
|
if od_config.model_config:
|
|
config = Pi0Config.from_model_config(dict(od_config.model_config))
|
|
# The deploy yaml is authoritative, but must not silently drop
|
|
# checkpoint-derived fields it omits. Backfill the camera order and
|
|
# any normalization stats the yaml didn't specify.
|
|
if self.model_dir and (not config.image_feature_keys or config.norm_stats is None):
|
|
ckpt = Pi0Config.from_pretrained(self.model_dir)
|
|
if not config.image_feature_keys:
|
|
config.image_feature_keys = ckpt.image_feature_keys
|
|
if not config.input_features:
|
|
config.input_features = ckpt.input_features
|
|
if config.norm_stats is None:
|
|
config.norm_stats = ckpt.norm_stats
|
|
return config
|
|
if self.model_dir:
|
|
return Pi0Config.from_pretrained(self.model_dir)
|
|
return Pi0Config()
|
|
|
|
def _resolve_tokenizer_source(self) -> str:
|
|
"""Prefer the checkpoint dir if it ships tokenizer files; else PaliGemma."""
|
|
if self.model_dir and os.path.isdir(self.model_dir):
|
|
if os.path.exists(os.path.join(self.model_dir, "tokenizer_config.json")):
|
|
return self.model_dir
|
|
return DEFAULT_PI0_TOKENIZER
|
|
|
|
@staticmethod
|
|
def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype:
|
|
dt = od_config.dtype
|
|
if isinstance(dt, torch.dtype):
|
|
return dt
|
|
return getattr(torch, str(dt).split(".")[-1], torch.float32)
|
|
|
|
@staticmethod
|
|
def _resolve_device(od_config: OmniDiffusionConfig) -> torch.device:
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
|
|
try:
|
|
return get_local_device()
|
|
except Exception: # noqa: BLE001
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
def _load_tokenizer(self):
|
|
from transformers import AutoTokenizer
|
|
|
|
return AutoTokenizer.from_pretrained(self.tokenizer_source)
|
|
|
|
def has_real_checkpoint(self) -> bool:
|
|
return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors"))
|
|
|
|
def _initialize_model(self) -> Pi0ForActionPrediction:
|
|
model = Pi0ForActionPrediction(self.config)
|
|
if self.has_real_checkpoint():
|
|
self._load_checkpoint(model)
|
|
else:
|
|
logger.info("Pi0Pipeline: no model.safetensors under %s; using random init.", self.model_dir)
|
|
model.to(device=self._device, dtype=self._torch_dtype)
|
|
model.eval()
|
|
return model
|
|
|
|
def _load_checkpoint(self, model: Pi0ForActionPrediction) -> None:
|
|
import safetensors.torch
|
|
|
|
path = os.path.join(self.model_dir, "model.safetensors")
|
|
logger.info("Pi0Pipeline: loading π0 weights from %s", path)
|
|
state = safetensors.torch.load_file(path)
|
|
model.load_weights(list(state.items()))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Framework weight-loading hook
|
|
# ------------------------------------------------------------------
|
|
def load_weights(self, weights=()): # noqa: D401
|
|
"""No-op for the diffusion loader: π0 self-loads its checkpoint in
|
|
``__init__`` (the kernel's ``load_weights`` handles the LeRobot remaps).
|
|
We expose no ``weights_sources``, so the loader passes an empty iterator
|
|
here; returning ``None`` skips its strict unloaded-weights check.
|
|
"""
|
|
for _ in weights:
|
|
pass
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Inference
|
|
# ------------------------------------------------------------------
|
|
@torch.inference_mode()
|
|
def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
|
|
extra_args = getattr(req.sampling_params, "extra_args", None) or {}
|
|
robot_obs = extra_args.get("robot_obs")
|
|
|
|
if robot_obs is None:
|
|
# Dummy warmup path (no obs): return zeros so engine warmup/capture
|
|
# doesn't crash. Mirrors DreamZero's dummy-run handling.
|
|
first_prompt = req.prompts[0] if req.prompts else ""
|
|
prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
|
|
num_steps = getattr(req.sampling_params, "num_inference_steps", None)
|
|
if prompt == "dummy run" or num_steps == 1:
|
|
logger.info("Pi0Pipeline: dummy warmup request without robot_obs — returning zeros.")
|
|
return DiffusionOutput(
|
|
output={
|
|
"actions": np.zeros(
|
|
(self.config.chunk_size, self.config.max_action_dim),
|
|
dtype=np.float32,
|
|
)
|
|
},
|
|
)
|
|
return DiffusionOutput(
|
|
error="Pi0Pipeline.forward requires sampling_params.extra_args['robot_obs'].",
|
|
)
|
|
|
|
images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
|
|
robot_obs, self.config, self.tokenizer, self._device
|
|
)
|
|
|
|
# State normalization (identity for pi0_base) inside the model.
|
|
state = self.model._normalize_state(state)
|
|
|
|
noise = extra_args.get("noise")
|
|
if noise is not None and not isinstance(noise, torch.Tensor):
|
|
noise = torch.as_tensor(noise, dtype=torch.float32, device=self._device)
|
|
elif isinstance(noise, torch.Tensor):
|
|
noise = noise.to(device=self._device, dtype=torch.float32)
|
|
|
|
num_steps = extra_args.get("num_inference_steps")
|
|
|
|
actions = self.model.sample_actions(
|
|
images=images,
|
|
image_masks=image_masks,
|
|
lang_tokens=lang_tokens,
|
|
lang_masks=lang_masks,
|
|
state=state,
|
|
noise=noise,
|
|
num_steps=num_steps,
|
|
)
|
|
actions = self.model._unnormalize_actions(actions)
|
|
|
|
# (B=1, horizon, action_dim) → (horizon, action_dim) numpy for the wire.
|
|
actions_np = actions.squeeze(0).float().cpu().numpy()
|
|
|
|
# Note: post-processing is applied engine-side via the registry
|
|
# (_DIFFUSION_POST_PROCESS_FUNCS), so we don't attach it here (a local
|
|
# closure wouldn't survive the orchestrator's multiprocess pickling).
|
|
return DiffusionOutput(output={"actions": actions_np})
|