MaciejBalaNVandHongsheng Liu
05ab509269
Signed-off-by: Maciej Bala <mbala@nvidia.com> Signed-off-by: MaciejBalaNV <mbala@nvidia.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Cosmos3 policy single-stage topology for online OpenPI serving.
|
|
|
|
Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one
|
|
diffusion stage: robot observation -> action chunk. They are
|
|
registered in ``OMNI_PIPELINES`` because online OpenPI serving
|
|
(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy
|
|
yaml's ``model_config.policy_server_config`` reaches the websocket handshake.
|
|
|
|
This pipeline declares neither ``hf_architectures`` nor
|
|
``diffusers_class_name``: policy checkpoints share their HF metadata
|
|
(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name=
|
|
Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which
|
|
must keep resolving through the default single-stage diffusion fallback.
|
|
Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g.
|
|
``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config
|
|
/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``.
|
|
"""
|
|
|
|
from vllm_omni.config.stage_config import (
|
|
PipelineConfig,
|
|
StageExecutionType,
|
|
StagePipelineConfig,
|
|
)
|
|
|
|
COSMOS3_POLICY_PIPELINE = PipelineConfig(
|
|
model_type="cosmos3_policy",
|
|
model_arch="Cosmos3OmniDiffusersPipeline",
|
|
stages=(
|
|
StagePipelineConfig(
|
|
stage_id=0,
|
|
model_stage="diffusion",
|
|
execution_type=StageExecutionType.DIFFUSION,
|
|
input_sources=(),
|
|
final_output=True,
|
|
final_output_type="action",
|
|
model_arch="Cosmos3OmniDiffusersPipeline",
|
|
),
|
|
),
|
|
)
|