vllm_omni.diffusion.models.pi0 ¶

Encode images with SigLIP vision tower + PaliGemma projector.

We run the two steps explicitly instead of calling PaliGemmaModel.get_image_features because that convenience method historically has returned either the projected tensor (old) or the raw vision-tower output (new). Being explicit makes us independent of which transformers release we're on.

Embed language tokens with PaliGemma's embedding table, returning the * sqrt(hidden)-scaled embedding (PaliGemma/Gemma convention).

This makes embed_language_tokens return the canonical scaled embedding on every transformers version.

Dispatch to prefix_only / suffix_only and return ([prefix_out, suffix_out], past_key_values_or_None).

Build from a config dict (LeRobot config.json or deploy yaml).

Keeps only the recognized dataclass fields (the LeRobot config also carries many training-only keys) and coerces image_resolution to a tuple.

Build from a checkpoint directory's config.json.

Apply one flow-matching denoising step: predict v_t from x_t.

Uses the prefix KV cache from sample_actions and only runs the action expert. Ref: openpi PI0Pytorch.denoise_step.

Build the prefix embeddings, per-token padding mask, and AR mask.

Prefix tokens form a contiguous sequence [img_cam_0..., img_cam_1..., ..., lang_tokens...] with bidirectional attention; the returned att_masks are all zeros because each token is free to attend to every other prefix token (the suffix pass will put a causal boundary right before the state token).

Ref: openpi PI0Pytorch.embed_prefix

Build the suffix embeddings + masks: [state_token, action_tokens×H].

AR mask layout: [1, 1, 0, 0, ..., 0] — the state token is a causal boundary (no later token attends backwards through it onto prefix by mistake), the first action token starts a new causal block, and the rest of the action tokens attend to each other bidirectionally.

Ref: openpi PI0Pytorch.embed_suffix

Load weights from an lerobot/pi0_base safetensors checkpoint.

Checkpoint keys (after stripping the leading model. prefix) look like paligemma_with_expert.paligemma.<sub>.*. Modern transformers nests those sub-modules one level deeper (paligemma.model.<sub>), and PaliGemma ties lm_head.weight with embed_tokens.weight at post_init time. Two remap rules make the load lossless:

1. paligemma.{vision_tower,multi_modal_projector,language_model}.* → paligemma.model.{...}.*
2. paligemma.lm_head.weight → paligemma.model.language_model.embed_tokens.weight (the checkpoint does not store embed_tokens.weight on its own — only the tied lm_head.weight copy — and PaliGemmaForConditionalGeneration does not register lm_head.weight as a Parameter when tied, so without this rule the language embedding would silently remain at its random init.)

On any remaining mismatch the loader logs a warning listing a sample of the unmatched checkpoint keys and model params; downstream parity tests surface these as a hard failure.

Generate an action chunk via iterative flow-matching denoising.

Convention: t=1 is noise, t=0 is the target — opposite of the published π0 paper but matches both OpenPI and LeRobot. Ref: openpi PI0Pytorch.sample_actions

Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot.

Accept a PIL image, an HWC uint8/float ndarray, or a CHW tensor and return a (1, 3, image_size, image_size) float tensor in [-1, 1].

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.