72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""π0.5 VLA model for vllm-omni.
|
|
|
|
PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert with AdaRMS
|
|
timestep conditioning + a flow-matching action head. Outputs a continuous action
|
|
chunk ``[horizon, action_dim]`` rather than tokens.
|
|
|
|
Differs from π0 in five places: AdaRMS timestep conditioning (``time_mlp_*``
|
|
instead of ``action_time_mlp_*``), AdaRMS norms throughout the action expert, a
|
|
200-token tokenizer budget, a discretized state carried in the prompt (no
|
|
``state_proj``), and relative-action support.
|
|
"""
|
|
|
|
from vllm_omni.diffusion.models.pi05.config import (
|
|
Pi05Config,
|
|
UnsupportedCheckpointCapabilityError,
|
|
resolve_excluded_action_indices,
|
|
)
|
|
from vllm_omni.diffusion.models.pi05.modeling_pi05 import (
|
|
GemmaVariantConfig,
|
|
PaliGemmaWithActionExpertPi05,
|
|
Pi05AdaRMSNorm,
|
|
Pi05ForActionPrediction,
|
|
create_sinusoidal_pos_embedding,
|
|
get_gemma_config,
|
|
make_att_2d_masks,
|
|
prepare_attention_masks_4d,
|
|
)
|
|
from vllm_omni.diffusion.models.pi05.pipeline_pi05 import (
|
|
Pi05Pipeline,
|
|
get_pi05_post_process_func,
|
|
)
|
|
from vllm_omni.diffusion.models.pi05.processor_pi05 import (
|
|
Pi05ImageProcessor,
|
|
Pi05Processor,
|
|
Pi05RelativeActions,
|
|
apply_norm,
|
|
build_norm_stats,
|
|
build_pi05_prompt,
|
|
discretize_state,
|
|
pil_image_to_tensor,
|
|
resize_with_pad,
|
|
tokenize_prompt,
|
|
)
|
|
|
|
__all__ = [
|
|
"Pi05Config",
|
|
"UnsupportedCheckpointCapabilityError",
|
|
"resolve_excluded_action_indices",
|
|
"Pi05ForActionPrediction",
|
|
"PaliGemmaWithActionExpertPi05",
|
|
"Pi05AdaRMSNorm",
|
|
"GemmaVariantConfig",
|
|
"get_gemma_config",
|
|
"create_sinusoidal_pos_embedding",
|
|
"make_att_2d_masks",
|
|
"prepare_attention_masks_4d",
|
|
"Pi05ImageProcessor",
|
|
"Pi05Processor",
|
|
"Pi05RelativeActions",
|
|
"build_pi05_prompt",
|
|
"discretize_state",
|
|
"apply_norm",
|
|
"build_norm_stats",
|
|
"pil_image_to_tensor",
|
|
"resize_with_pad",
|
|
"tokenize_prompt",
|
|
"Pi05Pipeline",
|
|
"get_pi05_post_process_func",
|
|
]
|