312 lines
13 KiB
Python
312 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""π0.5 VLA pipeline for vllm-omni.
|
|
|
|
Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the
|
|
DreamZero contract and the π0 pipeline: 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})``.
|
|
|
|
π0.5 is stateless across calls (no KV reuse, first-order Markov), so
|
|
``session_id`` / ``reset`` from the OpenPI protocol are accepted but ignored.
|
|
|
|
The post-processing order is load-bearing and matches LeRobot::
|
|
|
|
unnormalize → absolute actions → to_cpu
|
|
|
|
``AbsoluteActionsProcessorStep`` must run *after* unnormalization, because a
|
|
relative-action checkpoint's ``norm_stats`` are computed in relative space.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import fields as dataclass_fields
|
|
|
|
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.pi05.config import SUPPORTED_DTYPE_NAMES, Pi05Config
|
|
from vllm_omni.diffusion.models.pi05.modeling_pi05 import Pi05ForActionPrediction
|
|
from vllm_omni.diffusion.models.pi05.processor_pi05 import Pi05Processor
|
|
from vllm_omni.diffusion.models.pi05_pipeline_config import PI05_PIPELINE as PI05_PIPELINE
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# π0.5 pins the PaliGemma tokenizer (LeRobot hardcodes it too).
|
|
DEFAULT_PI05_TOKENIZER = "google/paligemma-3b-pt-224"
|
|
|
|
# float16 is absent deliberately, not by oversight: nothing here has been
|
|
# validated against it. The float32/bfloat16 trade-off is in the deploy config
|
|
# and recipes/lerobot/Pi05.md.
|
|
SUPPORTED_DTYPES = (torch.float32, torch.bfloat16)
|
|
|
|
# The two lists guard different entry points — what the checkpoint declares
|
|
# versus the dtype actually cast to — so they must not drift apart.
|
|
assert {str(dtype).split(".")[-1] for dtype in SUPPORTED_DTYPES} == set(SUPPORTED_DTYPE_NAMES)
|
|
|
|
|
|
def _checkpoint_declared_keys(model_dir: str) -> set[str]:
|
|
"""The keys the checkpoint's config.json actually carries."""
|
|
path = os.path.join(model_dir, "config.json")
|
|
if not os.path.exists(path):
|
|
return set()
|
|
with open(path, encoding="utf-8") as f:
|
|
return set(json.load(f))
|
|
|
|
|
|
def _comparable(value):
|
|
"""Compare yaml and checkpoint values without tripping on list/tuple."""
|
|
return list(value) if isinstance(value, (list, tuple)) else value
|
|
|
|
|
|
def _pi05_post_process(x):
|
|
"""Module-level identity post-process (picklable across the orchestrator's
|
|
multiprocess boundary — a local closure is not)."""
|
|
return x
|
|
|
|
|
|
def get_pi05_post_process_func(od_config: OmniDiffusionConfig):
|
|
"""π0.5 returns actions directly; post-processing is identity."""
|
|
del od_config
|
|
return _pi05_post_process
|
|
|
|
|
|
_LEROBOT_FLOAT32_IN_BFLOAT16 = (
|
|
"vision_tower",
|
|
"multi_modal_projector",
|
|
"input_layernorm",
|
|
"post_attention_layernorm",
|
|
"model.norm",
|
|
)
|
|
|
|
|
|
def _to_bfloat16_for_inference(model: Pi05ForActionPrediction) -> None:
|
|
"""Apply LeRobot's mixed-precision inference layout explicitly."""
|
|
inner_model = model.paligemma_with_expert
|
|
for name, param in inner_model.named_parameters():
|
|
target_dtype = (
|
|
torch.float32 if any(selector in name for selector in _LEROBOT_FLOAT32_IN_BFLOAT16) else torch.bfloat16
|
|
)
|
|
param.data = param.data.to(dtype=target_dtype)
|
|
for buffer in inner_model.buffers():
|
|
if buffer.is_floating_point():
|
|
buffer.data = buffer.data.to(dtype=torch.bfloat16)
|
|
|
|
for name, param in model.named_parameters():
|
|
if not name.startswith("paligemma_with_expert."):
|
|
param.data = param.data.to(dtype=torch.float32)
|
|
for name, buffer in model.named_buffers():
|
|
if not name.startswith("paligemma_with_expert.") and buffer.is_floating_point():
|
|
buffer.data = buffer.data.to(dtype=torch.float32)
|
|
|
|
|
|
def _set_inference_dtype(model: Pi05ForActionPrediction, dtype: torch.dtype) -> None:
|
|
if dtype == torch.float32:
|
|
model.to(dtype=torch.float32)
|
|
elif dtype == torch.bfloat16:
|
|
_to_bfloat16_for_inference(model)
|
|
else:
|
|
raise ValueError(f"Unsupported π0.5 inference dtype: {dtype!r}.")
|
|
|
|
|
|
class Pi05Pipeline(nn.Module):
|
|
"""π0.5 VLA pipeline: raw robot obs → continuous action chunk.
|
|
|
|
Registered as ``"Pi05Pipeline"`` in the diffusion registry.
|
|
"""
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
self.prefix = prefix
|
|
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()))
|
|
|
|
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()
|
|
|
|
self.processor = Pi05Processor(self.config, self.tokenizer, self._device)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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."""
|
|
if not model:
|
|
return None
|
|
if os.path.isdir(model):
|
|
return model
|
|
# Via repo_utils' shared HfApi rather than huggingface_hub directly, so
|
|
# the download carries vLLM-Omni's user agent like every other repo access.
|
|
from vllm_omni.transformers_utils.repo_utils import hf_api
|
|
|
|
return hf_api().snapshot_download(
|
|
repo_id=model,
|
|
allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"],
|
|
)
|
|
|
|
def _build_config(self, od_config: OmniDiffusionConfig) -> Pi05Config:
|
|
"""Read the config from the checkpoint, then let the deploy yaml override it."""
|
|
checkpoint_config = Pi05Config.from_pretrained(self.model_dir) if self.model_dir else None
|
|
if checkpoint_config is None:
|
|
return Pi05Config.from_model_config(od_config.model_config)
|
|
if not od_config.model_config:
|
|
return checkpoint_config
|
|
|
|
resolved = {
|
|
item.name: getattr(checkpoint_config, item.name) for item in dataclass_fields(Pi05Config) if item.init
|
|
}
|
|
# Only a key the checkpoint actually declares can disagree with the yaml.
|
|
# Comparing against the dataclass default instead would report every
|
|
# serving-only field, such as policy_server_config, on every start.
|
|
declared_keys = _checkpoint_declared_keys(self.model_dir)
|
|
for key, value in od_config.model_config.items():
|
|
if key in resolved and key in declared_keys and _comparable(value) != _comparable(resolved[key]):
|
|
logger.warning(
|
|
"Pi05Pipeline: the deploy config sets %s=%r, overriding %r from the checkpoint.",
|
|
key,
|
|
value,
|
|
resolved[key],
|
|
)
|
|
resolved.update(od_config.model_config)
|
|
return Pi05Config.from_model_config(resolved)
|
|
|
|
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_PI05_TOKENIZER
|
|
|
|
@staticmethod
|
|
def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype:
|
|
"""Resolve the dtype the weights are actually cast to.
|
|
|
|
This is the load-bearing check, not ``Pi05Config.dtype``: the cast in
|
|
:meth:`_initialize_model` reads the *top-level* ``OmniDiffusionConfig``
|
|
field, so a guard on the model config alone would let an unsupported
|
|
dtype through. See :data:`SUPPORTED_DTYPES` for why float16 is excluded.
|
|
"""
|
|
dt = od_config.dtype
|
|
resolved = dt if isinstance(dt, torch.dtype) else getattr(torch, str(dt).split(".")[-1], None)
|
|
if resolved not in SUPPORTED_DTYPES:
|
|
raise ValueError(
|
|
f"Unsupported π0.5 dtype: {dt!r}. Supported: "
|
|
f"{', '.join(sorted(str(d).split('.')[-1] for d in SUPPORTED_DTYPES))}."
|
|
)
|
|
return resolved
|
|
|
|
@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
|
|
|
|
# padding_side="right" is part of the π0.5 spec: the prefix is a fixed
|
|
# 200-token block whose live tokens must start at index 0.
|
|
return AutoTokenizer.from_pretrained(self.tokenizer_source, padding_side="right")
|
|
|
|
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) -> Pi05ForActionPrediction:
|
|
if not self.has_real_checkpoint():
|
|
expected = os.path.join(self.model_dir or "<missing-model-dir>", "model.safetensors")
|
|
raise FileNotFoundError(f"π0.5 serving requires checkpoint weights at {expected}.")
|
|
model = Pi05ForActionPrediction(self.config)
|
|
_set_inference_dtype(model, self._torch_dtype)
|
|
model.to(device=self._device)
|
|
self._load_checkpoint(model)
|
|
model.eval()
|
|
return model
|
|
|
|
def _load_checkpoint(self, model: Pi05ForActionPrediction) -> None:
|
|
import safetensors.torch
|
|
|
|
path = os.path.join(self.model_dir, "model.safetensors")
|
|
logger.info("Pi05Pipeline: loading π0.5 weights from %s", path)
|
|
try:
|
|
state = safetensors.torch.load_file(path)
|
|
model.load_weights(state.items())
|
|
except Exception as exc:
|
|
raise RuntimeError(f"Failed to load complete π0.5 checkpoint {path}: {exc}") from exc
|
|
|
|
# ------------------------------------------------------------------
|
|
# Framework weight-loading hook
|
|
# ------------------------------------------------------------------
|
|
def load_weights(self, weights=()): # noqa: D401
|
|
"""No-op for the diffusion loader: π0.5 self-loads its checkpoint."""
|
|
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("Pi05Pipeline: dummy warmup request without robot_obs — returning zeros.")
|
|
return DiffusionOutput(
|
|
output={
|
|
"actions": np.zeros(
|
|
(self.config.chunk_size, self.config.action_dim),
|
|
dtype=np.float32,
|
|
)
|
|
},
|
|
)
|
|
return DiffusionOutput(
|
|
error="Pi05Pipeline.forward requires sampling_params.extra_args['robot_obs'].",
|
|
)
|
|
|
|
images, image_masks, lang_tokens, lang_masks = self.processor.build_model_inputs(robot_obs)
|
|
|
|
num_steps = getattr(req.sampling_params, "num_inference_steps", None)
|
|
if num_steps is not None and (
|
|
isinstance(num_steps, bool) or not isinstance(num_steps, (int, np.integer)) or int(num_steps) < 1
|
|
):
|
|
raise ValueError(f"num_inference_steps must be a positive integer, got {num_steps!r}.")
|
|
if num_steps is not None:
|
|
num_steps = int(num_steps)
|
|
|
|
actions = self.model.sample_actions(
|
|
images=images,
|
|
image_masks=image_masks,
|
|
lang_tokens=lang_tokens,
|
|
lang_masks=lang_masks,
|
|
num_steps=num_steps,
|
|
generator=req.sampling_params.generator,
|
|
)
|
|
|
|
return DiffusionOutput(output={"actions": self.processor.build_model_outputs(actions, robot_obs)})
|