anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Support DROID policy server for Cosmos3 OpenPI (#4282)

Browse Source 
Signed-off-by: yuzhud <yuzhudong@gmail.com>
This commit is contained in:
Yuzhu Dong authored and GitHub committed 2026-06-15 21:22:59 +08:00
1 parent ae937320ff
commit e2e79177db
16 files changed
+1397 -154
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+163
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import sys
import types
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any
@@ -230,6 +231,8 @@ def test_pipeline_registered_and_exported() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.registry import (
_DIFFUSION_ACTION_POST_PROCESS_FUNCS,
_DIFFUSION_IR_OP_PRIORITY_FUNCS,
_DIFFUSION_MODELS,
_DIFFUSION_POST_PROCESS_FUNCS,
_DIFFUSION_PRE_PROCESS_FUNCS,
@@ -245,6 +248,10 @@ def test_pipeline_registered_and_exported() -> None:
)
assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
assert (
_DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func"
)
assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func"
assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
@@ -413,6 +420,73 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None:
func({"image": video, "video": video})
def test_action_postprocess_handles_robolab_policy_outputs() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
RoboLabPolicyInputs,
get_cosmos3_action_post_process_func,
make_robolab_action_postprocess_inputs,
)
func = get_cosmos3_action_post_process_func(SimpleNamespace())
inputs = RoboLabPolicyInputs(
prompt="Pick the cube.",
video_tensor=torch.zeros(1, 3, 3, 16, 16),
action_tensor=torch.zeros(2, 2),
action_condition_indexes=[0],
action_start_frame_offset=1,
raw_action_dim=2,
domain_id=7,
fps=15.0,
height=16,
width=16,
image_size=None,
num_frames=3,
num_inference_steps=4,
guidance_scale=3.0,
flow_shift=5.0,
seed=11,
history_length=1,
action_space="joint_pos",
observation={},
)
action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)}
processed = func(action, custom_output=custom_output)
assert processed.shape == (1, 2)
assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]]))
assert "robolab_action_postprocess" not in custom_output
def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_ir_op_priority_func
@dataclass
class FakeIrOpPriorityConfig:
rms_norm: list[str]
fused_add_rms_norm: list[str]
custom_op: list[str]
fake_kernel = types.ModuleType("vllm.config.kernel")
fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
func = get_cosmos3_ir_op_priority_func(SimpleNamespace())
default_priority = FakeIrOpPriorityConfig(
rms_norm=["vllm_c", "native"],
fused_add_rms_norm=["vllm_c", "native"],
custom_op=["platform_kernel", "native"],
)
merged = func(default_priority, vllm_config=SimpleNamespace())
assert merged.rms_norm == ["native"]
assert merged.fused_add_rms_norm == ["native"]
assert merged.custom_op == ["platform_kernel", "native"]
def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
@@ -762,6 +836,95 @@ class TestForwardRouting:
)
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
assert output.custom_output["action"].shape == (1, 2, 2)
assert "action_only_output" not in output.custom_output
def test_forward_dispatches_robolab_policy_flow(
self,
make_cosmos3_pipeline,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
pipeline = make_cosmos3_pipeline()
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
captured = self._install_forward_stubs(pipeline)
video_latents = torch.zeros(1, 2, 1, 1, 1)
velocity_mask = torch.ones(1, 1, 1, 1, 1)
condition_latents = torch.zeros_like(video_latents)
inputs = pipeline_cosmos3.RoboLabPolicyInputs(
prompt="Pick the cube.",
video_tensor=torch.zeros(1, 3, 3, 16, 16),
action_tensor=torch.zeros(2, 2),
action_condition_indexes=[0],
action_start_frame_offset=1,
raw_action_dim=2,
domain_id=7,
fps=15.0,
height=16,
width=16,
image_size=None,
num_frames=3,
num_inference_steps=4,
guidance_scale=3.0,
flow_shift=5.0,
seed=11,
history_length=1,
action_space="joint_pos",
observation={},
)
def fake_prepare_action_latents(**kwargs):
captured["prepare_action"] = kwargs
action_chunk_size = kwargs["action_chunk_size"]
raw_action_dim = int(kwargs["raw_action_dim"])
return (
torch.zeros(1, action_chunk_size, 4),
torch.ones(1, action_chunk_size, 1),
torch.zeros(1, action_chunk_size, 4),
raw_action_dim,
)
def fake_prepare_action_video(*args, **kwargs):
captured["prepare_action_video"] = {"args": args, "kwargs": kwargs}
return video_latents, velocity_mask, condition_latents
monkeypatch.setattr(
pipeline_cosmos3,
"build_robolab_unipc_scheduler",
lambda num_steps, shift, device: StubScheduler(list(range(num_steps, 0, -1)), flow_shift=shift),
)
pipeline._build_robolab_policy_inputs = lambda sp, prompt_data, request_id=None: inputs
pipeline._prepare_action_latents = fake_prepare_action_latents
pipeline._prepare_latents_action_video = fake_prepare_action_video
pipeline._decode_latents = lambda latents: (_ for _ in ()).throw(
AssertionError("RoboLab should not decode video")
)
output = pipeline.forward(SimpleNamespace(prompts=["ignored"], sampling_params=make_sampling_params()))
assert captured["format"] == {
"prompt": "Pick the cube.",
"negative_prompt": "",
"num_frames": 3,
"frame_rate": 15.0,
"height": 16,
"width": 16,
"is_t2i": False,
}
assert "flow_shifts" not in captured
assert pipeline.scheduler.set_timesteps_calls == []
assert captured["prepare_action"]["clean_action"] is inputs.action_tensor
assert captured["prepare_action"]["condition_indexes"] == [0]
assert captured["prepare_action_video"]["kwargs"] == {"image_size": None}
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1]
assert output.output == {}
assert output.custom_output["action_only_output"] is True
assert output.custom_output["action"].shape == (1, 2, 2)
assert "actions" not in output.custom_output
assert "robolab_action_postprocess" in output.custom_output
assert "robolab_policy_inputs" not in output.custom_output
@pytest.mark.parametrize(
("prompt", "sampling_params", "message"),
tests/diffusion/test_diffusion_plugin_hooks.py
+33 -1
View File
@@ -10,11 +10,14 @@ This module tests:
- Worker integration: model runner resolved via platform hook
"""
from unittest.mock import patch
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
from vllm_omni.diffusion.registry import (
_DIFFUSION_ACTION_POST_PROCESS_FUNCS,
_DIFFUSION_IR_OP_PRIORITY_FUNCS,
_DIFFUSION_MODELS,
_DIFFUSION_POST_PROCESS_FUNCS,
_DIFFUSION_PRE_PROCESS_FUNCS,
@@ -59,6 +62,8 @@ class TestRegisterDiffusionModel:
original_models = _DIFFUSION_MODELS.copy()
original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
yield
_DIFFUSION_MODELS.clear()
_DIFFUSION_MODELS.update(original_models)
@@ -66,6 +71,10 @@ class TestRegisterDiffusionModel:
_DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
_DIFFUSION_POST_PROCESS_FUNCS.clear()
_DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
_DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
_DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
_DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
_DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
def test_register_new_model(self):
"""Test registering a new diffusion model with pre/post process functions."""
@@ -75,6 +84,8 @@ class TestRegisterDiffusionModel:
class_name="TestPipeline",
pre_process_func_name="test_pre_process",
post_process_func_name="test_post_process",
action_post_process_func_name="test_action_post_process",
ir_op_priority_func_name="test_ir_op_priority",
)
assert "TestPipeline" in _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -84,6 +95,8 @@ class TestRegisterDiffusionModel:
)
assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
class TestWorkerUsesHook:
@@ -108,3 +121,22 @@ class TestWorkerUsesHook:
assert worker.model_runner is mock_runner_instance
mock_platform.get_diffusion_model_runner_cls.assert_called_once()
mock_resolve.assert_called_once_with("custom.path")
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_diffusion_ir_op_priority_func")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
def test_ir_op_priority_hook_receives_platform_default(self, mock_platform, mock_get_hook):
"""Test model IR priority hook merges from the platform default."""
from vllm_omni.diffusion.worker.diffusion_worker import _resolve_ir_op_priority
od_config = SimpleNamespace(model_class_name="TestPipeline")
vllm_config = SimpleNamespace()
default_priority = object()
merged_priority = object()
hook = Mock(return_value=merged_priority)
mock_platform.get_default_ir_op_priority.return_value = default_priority
mock_get_hook.return_value = hook
assert _resolve_ir_op_priority(od_config, vllm_config) is merged_priority
mock_platform.get_default_ir_op_priority.assert_called_once_with(vllm_config)
mock_get_hook.assert_called_once_with(od_config)
hook.assert_called_once_with(default_priority, vllm_config=vllm_config)
tests/diffusion/test_diffusion_scheduler.py
+72
View File
@@ -567,6 +567,78 @@ class TestDiffusionEngine:
torch.tensor([3.0, 4.0]),
)
@pytest.mark.asyncio
async def test_step_empty_dict_output_still_runs_postprocess(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(
model_class_name="mock_model",
enable_cpu_offload=False,
)
engine.pre_process_func = None
engine.post_process_func = mocker.Mock(return_value={"video": ["processed"]})
engine._post_process_accepts_sampling_params = False
engine._check_and_start_background_loop = mocker.AsyncMock()
engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
return_value=DiffusionOutput(
output={},
custom_output={"actions": torch.tensor([[1.0, 2.0]])},
)
)
request = OmniDiffusionRequest(
prompts=["prompt"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="req-action",
)
mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
outputs = await engine.step(request)
engine.post_process_func.assert_called_once_with({})
assert outputs[0].images == ["processed"]
torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[1.0, 2.0]]))
@pytest.mark.asyncio
async def test_step_action_only_flag_skips_postprocess(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(
model_class_name="mock_model",
enable_cpu_offload=False,
)
engine.pre_process_func = None
engine.post_process_func = mocker.Mock(side_effect=AssertionError("postprocess should be skipped"))
engine.action_post_process_func = mocker.Mock(return_value=torch.tensor([[3.0, 4.0]]))
engine._post_process_accepts_sampling_params = False
engine._action_post_process_accepts_custom_output = True
engine._action_post_process_accepts_sampling_params = False
engine._check_and_start_background_loop = mocker.AsyncMock()
raw_action = torch.tensor([[1.0, 2.0]])
engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
return_value=DiffusionOutput(
output={},
custom_output={
"action": raw_action,
"action_only_output": True,
},
)
)
request = OmniDiffusionRequest(
prompts=["prompt"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="req-action",
)
mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
outputs = await engine.step(request)
engine.post_process_func.assert_not_called()
engine.action_post_process_func.assert_called_once()
assert engine.action_post_process_func.call_args.args[0] is raw_action
assert "custom_output" in engine.action_post_process_func.call_args.kwargs
assert outputs[0].images == []
torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[3.0, 4.0]]))
class TestStepScheduler:
def setup_method(self) -> None: