Yichen WangandDidan Deng
2a9f6aaa8e
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""π0 (Pi-Zero) VLA model for vllm-omni.
|
|
|
|
PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert + a
|
|
flow-matching action head. Outputs a continuous action chunk
|
|
``[horizon, action_dim]`` rather than tokens.
|
|
"""
|
|
|
|
from vllm_omni.diffusion.models.pi0.config import Pi0Config
|
|
from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
|
|
GemmaVariantConfig,
|
|
PaliGemmaWithActionExpert,
|
|
Pi0ForActionPrediction,
|
|
create_sinusoidal_pos_embedding,
|
|
get_gemma_config,
|
|
make_att_2d_masks,
|
|
prepare_attention_masks_4d,
|
|
)
|
|
from vllm_omni.diffusion.models.pi0.pipeline_pi0 import (
|
|
Pi0Pipeline,
|
|
get_pi0_post_process_func,
|
|
)
|
|
from vllm_omni.diffusion.models.pi0.processor_pi0 import (
|
|
Pi0ImageProcessor,
|
|
build_model_inputs,
|
|
pil_image_to_tensor,
|
|
resize_with_pad,
|
|
tokenize_prompt,
|
|
)
|
|
|
|
__all__ = [
|
|
"Pi0Config",
|
|
"Pi0ForActionPrediction",
|
|
"PaliGemmaWithActionExpert",
|
|
"GemmaVariantConfig",
|
|
"get_gemma_config",
|
|
"create_sinusoidal_pos_embedding",
|
|
"make_att_2d_masks",
|
|
"prepare_attention_masks_4d",
|
|
"Pi0ImageProcessor",
|
|
"build_model_inputs",
|
|
"pil_image_to_tensor",
|
|
"resize_with_pad",
|
|
"tokenize_prompt",
|
|
"Pi0Pipeline",
|
|
"get_pi0_post_process_func",
|
|
]
|