Skip to content

vllm_omni.diffusion.models.pi0.pipeline_pi0

π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.

DEFAULT_PI0_TOKENIZER module-attribute 

DEFAULT_PI0_TOKENIZER = 'google/paligemma-3b-pt-224'

logger module-attribute 

logger = init_logger(__name__)

Pi0Pipeline

Bases: 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).

config instance-attribute 

config = self._build_config(od_config)

model instance-attribute 

model = self._initialize_model()

model_dir instance-attribute 

model_dir = self._resolve_model_dir(od_config.model)

od_config instance-attribute 

od_config = od_config

prefix instance-attribute 

prefix = prefix

tokenizer instance-attribute 

tokenizer = self._load_tokenizer()

tokenizer_source instance-attribute 

tokenizer_source = str(
    custom_args.get(
        "tokenizer", self._resolve_tokenizer_source()
    )
)

forward 

forward(
    req: OmniDiffusionRequest, **kwargs
) -> DiffusionOutput

has_real_checkpoint 

has_real_checkpoint() -> bool

load_weights 

load_weights(weights=())

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.

get_pi0_post_process_func 

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.