[Diffusion] DreamZero world model integration with CFG parallel + OpenPI serving

[TKONIY]

Summary

Current PR branch: feature/dreamzero-pipeline
Latest pushed commit: a820f779 tests: fix diffusion scheduler mock imports after rebase

This PR integrates DreamZero into vllm-omni with:

• DreamZero diffusion pipeline support
• OpenPI-compatible robot WebSocket serving
• DreamZero-specific transform pipeline owned by the model
• DreamZero stage config + model-specific policy server config
• Online serving example with bundled real videos
• DROID sim-evals rollout client for the vLLM OpenPI endpoint
• Explicit optional-dependency errors in the DreamZero example clients
• Self-contained e2e tests plus optional upstream parity checks

What Changed

Model / Pipeline

• Added DreamZero stage config: vllm_omni/model_executor/stage_configs/dreamzero.yaml
• Added DreamZero model defaults in vllm_omni/diffusion/models/dreamzero/utils.py
• Moved robot transforms into vllm_omni/diffusion/models/dreamzero/transform/
• Let DreamZeroPipeline consume raw robot_obs and apply transform inside the model
• Kept reset semantics as deferred engine reset via extra_args["reset"]

Serving

• Added model-specific PolicyServerConfig loading in vllm_omni/entrypoints/openai/realtime/robot/openpi_serving.py
• Added optional dependency guard, structured errors, and idle timeout in vllm_omni/entrypoints/openai/realtime/robot/openpi_connection.py
• Added create_policy_server() so OpenPI serving is enabled only when the loaded model provides policy_server_config
• Sent OpenPI handshake metadata from model config instead of hardcoding it in the WebSocket connection layer

Framework Wiring

• vllm_omni/diffusion/stage_diffusion_proc.py
  • no longer carries DreamZero-specific inline detection
  • preserves explicit model_class_name for unregistered architectures
• vllm_omni/diffusion/diffusion_engine.py
  • deduplicates multimodal payload slicing for audio / actions
• vllm_omni/entrypoints/utils.py
  • resolves DreamZero config from model type override

Tests / Examples / Docs

• Added OpenPI unit tests:
  • tests/entrypoints/openai_api/test_openpi_connection.py
  • tests/entrypoints/openai_api/test_openpi_serving.py
• Added DreamZero e2e / example tests:
  • tests/e2e/online_serving/test_dreamzero.py
  • tests/examples/online_serving/test_dreamzero.py
• Added optional upstream parity checks under:
  • tests/dreamzero/upstream/
• Added runnable example:
  • examples/online_serving/dreamzero/
• Added offline prediction-video export helpers:
  • examples/online_serving/dreamzero/export_prediction_video.py
  • examples/online_serving/dreamzero/generate_comparison_videos.py
• Added DROID sim rollout client:
  • examples/online_serving/dreamzero/droid_sim_eval_client.py
• Added model docs:
  • docs/models/dreamzero/README.md
  • docs/models/dreamzero/quick_start.md
• Documented per-script environment requirements for:
  • vllm serve
  • bundled OpenPI client
  • prediction-video export helpers
  • DROID sim-eval client
  • optional upstream parity tests

Quick Start

Start the server

From the repository root:

CUDA_VISIBLE_DEVICES=0,1 \
examples/online_serving/dreamzero/run_server.sh

If you only want 1 GPU:

CUDA_VISIBLE_DEVICES=0 \
CFG_PARALLEL_SIZE=1 \
examples/online_serving/dreamzero/run_server.sh

OpenPI websocket endpoint:

• ws://127.0.0.1:8000/v1/realtime/robot/openpi

Run the client

From the repository root:

python examples/online_serving/dreamzero/openpi_client.py \
  --host 127.0.0.1 \
  --port 8000

Extra client dependencies:

pip install openpi-client websockets opencv-python

The example client uses bundled real videos from:

• examples/online_serving/dreamzero/assets/

Optional flags:

python examples/online_serving/dreamzero/openpi_client.py \
  --host 127.0.0.1 \
  --port 8000 \
  --video-dir examples/online_serving/dreamzero/assets \
  --session-id demo-session \
  --num-chunks 2

Run DROID sim-eval

Run this from an environment where isaaclab, isaaclab_tasks,
sim_evals, and gymnasium are already importable.

pip install openpi-client websockets opencv-python mediapy

From the vllm-omni repository root, invoke the client through an external
Isaac Lab launcher, for example:

/path/to/isaaclab.sh -p \
  examples/online_serving/dreamzero/droid_sim_eval_client.py \
  --host 127.0.0.1 \
  --port 8000 \
  --scene 1 \
  --episodes 1 \
  --headless \
  --device cuda:0

Validation

Passed locally:

• PYTHONPATH=. pytest tests/entrypoints/openai_api/test_openpi_serving.py tests/entrypoints/openai_api/test_openpi_connection.py -q
• OPENPI_E2E_GPUS=0,1 PYTHONPATH=. pytest tests/e2e/online_serving/test_dreamzero.py -q --run-level=advanced_model
• PYTHONPATH=. .venv/bin/python -m py_compile examples/online_serving/dreamzero/openpi_client.py examples/online_serving/dreamzero/droid_sim_eval_client.py

This confirms:

• DreamZero OpenPI handshake / serving works
• DreamZero online e2e works with real bundled videos

Performance Snapshot

Measurement scope for the numbers below:

• eager mode (--enforce-eager)
• no torch.compile
• no DiT cache / dynamic cached schedule in the baseline table
• single-request OpenPI serving path
• DreamZero parallelism configured via stage YAML parallel_config, not CLI TP/CFG flags

Hardware environment for these measurements:

• GPU: 4x NVIDIA RTX PRO 6000 Blackwell Server Edition, 97887 MiB VRAM each
• GPU driver: 590.48.01
• CPU: 2x AMD EPYC 9355 32-Core Processor (128 logical CPUs total)
• Host RAM: 1.5 TiB

VRAM

Interpretation notes:

• dreamzero.yaml is the default TP=1, CFG=1 baseline
• all GPUs on this host are the same model, so the table reports only how many GPUs were used and the per-GPU memory numbers
• TP=2, CFG=2 needs 4 GPUs and is still blocked on this host

Mode	GPUs Used	Per-GPU startup VRAM	Per-GPU peak VRAM	Status	Notes
TP=1, CFG=1	1	43.58 GiB	52.01 GB reserved	Measured	vllm_omni/model_executor/stage_configs/dreamzero.yaml
TP=1, CFG=2	2	43.58 GiB	49.61 GB reserved	Measured	True CFG-parallel serving via dreamzero_tp1_cfg2.yaml
TP=2, CFG=1	2	28.88 GiB	32.65 GB reserved	Measured	True TP serving via dreamzero_tp2_cfg1.yaml
TP=2, CFG=2	4	Not measured	Not measured	Pending	Requires 4 GPUs; blocked by unrelated jobs on GPUs 2,3

Latency

Interpretation notes:

• latency metric here is server-side DiffusionEngine.step breakdown ... total=... ms
• each row below used the same OpenPI client workload:
  • 3 action-producing requests (infer, infer, reset, infer)
  • same bundled DreamZero example videos
• TP=1, CFG=2 reduces wall time versus baseline on this workload
• TP=2, CFG=1 works correctly but is slower than baseline on this workload

Mode	GPUs Used	Mean latency	P50 latency	Range	Client wall time	Status	Notes
TP=1, CFG=1	1	7349.93 ms	7419.06 ms	7145.32–7485.42 ms	22.329 s	Measured	tmp/dreamzero_perf_yaml/tp1_cfg1.log
TP=1, CFG=2	2	4033.93 ms	3863.00 ms	3829.45–4409.35 ms	12.365 s	Measured	tmp/dreamzero_perf_yaml/tp1_cfg2.log
TP=2, CFG=1	2	8636.77 ms	8526.22 ms	8451.61–8932.49 ms	26.196 s	Measured	tmp/dreamzero_perf_yaml/tp2_cfg1.log
TP=2, CFG=2	4	Not measured	Not measured	Not measured	Not measured	Pending	Requires 4 GPUs; blocked by unrelated jobs on GPUs 2,3

Important Future Work

• Design clear and stable API and protocols for robot. Current design of PolicyServerConfig and Transform is ugly.

• Manage KV Cache with Paged Attention.
• Performance Optimizations, e.g., asynchronous pipelined execution.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[linyueqian]

Nice clean architecture with good separation of concerns (connection / serving / transform). A few items to address before merge.

[hsliuustc0106]

BLOCKING:

• Merge conflict — This PR is in CONFLICTING state. Please rebase onto latest main before review can proceed.

Non-blocking notes:

1. PR description TODO says "Clear the comments" — please resolve before merge.
2. Consider adding concrete latency / VRAM numbers to the PR description (even rough figures from the local validation runs). The current "passed locally" section lists test commands but not their measured outputs.

[TKONIY]

@yinpeiqi @hsliuustc0106 @linyueqian
There are a lot of unit tests under tests/dreamzero/upstream and end-to-end test to check if our implementation aligns with upstream DreamZero's official implementation. They require dependency to that repo. Do you prefer keeping or removing them?

[TKONIY]

BLOCKING:

• Merge conflict — This PR is in CONFLICTING state. Please rebase onto latest main before review can proceed.

Non-blocking notes:

1. PR description TODO says "Clear the comments" — please resolve before merge.
2. Consider adding concrete latency / VRAM numbers to the PR description (even rough figures from the local validation runs). The current "passed locally" section lists test commands but not their measured outputs.

All comments have been addressed.