[BugFix][Cosmos3] Register opt-in T2I pipeline so deploy-config applies (#7971)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

recipes/cosmos3/Cosmos3-Nano.md

@@ -126,7 +126,10 @@ speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
 `0` removes the streak cap.
 
 To run **without** guardrails (you are responsible for license compliance),
-add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
+add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
+For extra GPUs use
 `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
 `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
 `--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
@@ -728,8 +731,10 @@ vllm serve nvidia/Cosmos3-Nano \
 ```
 
 To run **without** guardrails (you are responsible for license compliance),
-add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
-add `--tensor-parallel-size 8`. `--quantization fp8` and
+add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`).
+For tensor parallel add `--tensor-parallel-size 8`. `--quantization fp8` and
 `--enable-layerwise-offload` are not supported on NPU.
 The pipeline auto-resolves from `model_index.json`; pass
 `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.

recipes/cosmos3/Cosmos3-Super.md

@@ -57,7 +57,13 @@ vllm serve nvidia/Cosmos3-Super \
 
 Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
 cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
-disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
+disable, or pass
+[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml)
+via `--deploy-config` on either multi-GPU command above (opt-in omni deploy
+overlay, `pipeline: cosmos3_omni_deploy`; without that `pipeline:` key
+`--deploy-config` cannot apply for Cosmos3). Same YAML works for Super / Nano
+(T2I and video); topology `final_output_type` matches the CLI default (`video`).
+`--enable-layerwise-offload` reduces VRAM on smaller GPUs;
 `--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the
 decoder convolutions to channels-last kernels (output no longer bit-identical to
 diffusers; the default `lossless` fast path is bit-exact, see

tests/config/test_config_factory.py

@@ -694,6 +694,105 @@ class TestCosmos3PolicyPipeline:
         assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position"
 
 
+class TestCosmos3OmniDeployPipeline:
+    """Omni deploy yaml selects an opt-in topology so --deploy-config applies (#6874)."""
+
+    def test_registered_without_auto_capturing_cosmos3_omni(self):
+        assert "cosmos3_omni_deploy" in OMNI_PIPELINES
+        # Same shared HF metadata as policy: must not register cosmos3_omni for
+        # auto-detect, or T2I/video/policy would collide.
+        assert "cosmos3_omni" not in OMNI_PIPELINES
+        pipeline = OMNI_PIPELINES["cosmos3_omni_deploy"]
+        assert pipeline.hf_architectures == ()
+        assert pipeline.diffusers_class_name is None
+        # Align with CLI Cosmos3OmniDiffusersPipeline default (video), not T2I-only.
+        assert pipeline.stages[0].final_output_type == "video"
+
+    def test_deploy_yaml_applies_guardrails_false(self):
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
+        assert deploy.pipeline == "cosmos3_omni_deploy"
+        assert deploy.stages[0].devices is None
+
+        stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_omni_deploy"], deploy)
+        assert len(stages) == 1
+        stage = stages[0].to_omegaconf()
+
+        assert stage.stage_type == "diffusion"
+        assert stage.final_output_type == "video"
+        assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline"
+        assert stage.engine_args.model_config.guardrails is False
+        assert "devices" not in stage.runtime
+
+    def test_deploy_yaml_accepts_documented_multi_gpu_cli(self):
+        from vllm_omni.engine.stage_init_utils import _check_stage_device_layout
+
+        deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml"))
+        stages, _ = StageConfigFactory._create_legacy_from_registry(
+            OMNI_PIPELINES["cosmos3_omni_deploy"],
+            {
+                "cfg_parallel_size": 2,
+                "use_hsdp": True,
+                "hsdp_shard_size": 2,
+            },
+            user_deploy_config=deploy,
+        )
+        assert len(stages) == 1
+        omega = stages[0].to_omegaconf()
+        assert omega.engine_args.model_config.guardrails is False
+        assert omega.engine_args.parallel_config.cfg_parallel_size == 2
+        assert omega.engine_args.parallel_config.use_hsdp is True
+        assert omega.engine_args.parallel_config.hsdp_shard_size == 2
+        assert "devices" not in omega.runtime
+        _check_stage_device_layout(omega, dict(omega.engine_args))
+
+    def test_get_pipeline_config_honours_deploy_yaml_pipeline_key(self):
+        deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
+        # No HF download: explicit pipeline: key is highest priority.
+        pipeline = StageConfigFactory.get_pipeline_config(
+            model="nvidia/Cosmos3-Super-Text2Image",
+            trust_remote_code=True,
+            deploy_config_path=deploy_path,
+        )
+        assert pipeline is not None
+        assert pipeline.model_type == "cosmos3_omni_deploy"
+
+    def test_without_deploy_yaml_stays_on_unregistered_fallback(self):
+        # Shared HF metadata must not auto-select a pipeline; only the deploy
+        # yaml ``pipeline:`` key selects cosmos3_omni_deploy and merges guardrails.
+        assert "cosmos3_omni" not in OMNI_PIPELINES
+
+        class FakeCosmos3Config(PretrainedConfig):
+            model_type = "cosmos3_omni"
+
+        fake_config = FakeCosmos3Config()
+        fake_config.architectures = ["Cosmos3OmniDiffusersPipeline"]
+        model = "nvidia/Cosmos3-Super-Text2Image"
+        deploy_path = get_deploy_config_path("cosmos3_omni.yaml")
+
+        StageConfigFactory.get_hf_config.cache_clear()
+        StageConfigFactory.try_infer_model_type.cache_clear()
+        with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config):
+            assert (
+                StageConfigFactory.get_pipeline_config(
+                    model=model,
+                    trust_remote_code=True,
+                    deploy_config_path=None,
+                )
+                is None
+            )
+
+            pipeline = StageConfigFactory.get_pipeline_config(
+                model=model,
+                trust_remote_code=True,
+                deploy_config_path=deploy_path,
+            )
+        assert pipeline is not None
+        assert pipeline.model_type == "cosmos3_omni_deploy"
+
+        stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))
+        assert stages[0].to_omegaconf().engine_args.model_config.guardrails is False
+
+
 class TestStagePipelineConfig:
     def test_frozen(self):
         s = StagePipelineConfig(stage_id=0, model_stage="a")

vllm_omni/config/pipeline_registry.py

@@ -49,7 +49,10 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
     BAGEL_THINK_PIPELINE,
 )
 from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE
-from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE
+from vllm_omni.model_executor.models.cosmos3.pipeline import (
+    COSMOS3_OMNI_DEPLOY_PIPELINE,
+    COSMOS3_POLICY_PIPELINE,
+)
 from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
@@ -154,11 +157,12 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "Gr00tN1d7": GR00T_N1D7_PIPELINE,
     "pi0": PI0_PIPELINE,
     "pi05": PI05_PIPELINE,
-    # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3
-    # checkpoints (which stay on the single-stage diffusion fallback), so this
-    # entry is only reachable through a deploy yaml's ``pipeline:`` key
-    # (see deploy/cosmos3_policy_droid.yaml).
+    # Cosmos3 policy / omni-deploy topologies share HF metadata with video
+    # Cosmos3 checkpoints (which stay on the single-stage diffusion fallback),
+    # so these entries are only reachable through a deploy yaml's ``pipeline:``
+    # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml).
     "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
+    "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE,
     "gepard": GEPARD_PIPELINE,
     "glm_image": GLM_IMAGE_PIPELINE,
     "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,

vllm_omni/deploy/cosmos3_omni.yaml

@@ -0,0 +1,30 @@
+# Cosmos3 omni deploy overlay: single co-located diffusion stage (both MoT
+# towers in one process). Topology (COSMOS3_OMNI_DEPLOY_PIPELINE) is declared in
+# vllm_omni/model_executor/models/cosmos3/pipeline.py.
+#
+# T2I / video / policy Cosmos3 checkpoints share HF metadata
+# (model_type=cosmos3_omni, model_index.json _class_name=
+# Cosmos3OmniDiffusersPipeline). This pipeline cannot be auto-detected and must
+# be selected explicitly so --deploy-config is applied (#6874):
+#
+#   vllm serve nvidia/Cosmos3-Super --omni \
+#     --deploy-config /absolute/path/to/vllm_omni/deploy/cosmos3_omni.yaml
+#
+# Same overlay works for Nano (including video). Without --deploy-config,
+# behaviour is unchanged (CLI single-stage fallback). --no-guardrails remains
+# the CLI-only equivalent of model_config.guardrails below. Per-checkpoint
+# knobs can still be layered with --stage-overrides.
+
+pipeline: cosmos3_omni_deploy
+# async_chunk defaults to True, which the deploy validator rejects for a
+# single-stage pipeline (no next-stage processor). Must be false here.
+async_chunk: false
+distributed_executor_backend: mp
+
+stages:
+  - stage_id: 0
+    model_class_name: Cosmos3OmniDiffusersPipeline
+    model_config:
+      # Gated nvidia/Cosmos-1.0-Guardrail is often unavailable; default this
+      # recommended layout off. Equivalent to launching with --no-guardrails.
+      guardrails: false

vllm_omni/model_executor/models/cosmos3/pipeline.py

@@ -1,21 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Cosmos3 policy single-stage topology for online OpenPI serving.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cosmos3 opt-in topologies for policy (OpenPI) and omni deploy YAML.
 
-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.
+Both topologies declare neither ``hf_architectures`` nor
+``diffusers_class_name``. Every Cosmos3 checkpoint (T2I, video, policy)
+shares HF metadata (``model_type=cosmos3_omni``, ``model_index.json``
+``_class_name=Cosmos3OmniDiffusersPipeline``). Auto-registering any of them
+under that metadata would capture the others. Select explicitly via a deploy
+yaml ``pipeline:`` key.
 
-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``.
+Policy (``cosmos3_policy``)::
+
+    vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \\
+      --deploy-config .../vllm_omni/deploy/cosmos3_policy_droid.yaml
+
+Omni deploy overlay (``cosmos3_omni_deploy``) — makes ``--deploy-config``
+apply stage / ``model_config`` (e.g. ``guardrails: false``) instead of being
+silently dropped (#6874). Same overlay works for Super / Nano (T2I and video).
+``final_output_type`` matches the CLI Cosmos3OmniDiffusersPipeline default
+(``video``)::
+
+    vllm serve nvidia/Cosmos3-Super --omni \\
+      --deploy-config .../vllm_omni/deploy/cosmos3_omni.yaml
+
+Without ``--deploy-config``, T2I/video keep the default single-stage
+diffusion CLI fallback. ``--no-guardrails`` remains the CLI-only path.
 """
 
 from vllm_omni.config.stage_config import (
@@ -39,3 +48,19 @@ COSMOS3_POLICY_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+COSMOS3_OMNI_DEPLOY_PIPELINE = PipelineConfig(
+    model_type="cosmos3_omni_deploy",
+    model_arch="Cosmos3OmniDiffusersPipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="video",
+            model_arch="Cosmos3OmniDiffusersPipeline",
+        ),
+    ),
+)