[Edge][WP6] Slim the AR/TTS import path and add an edge dependency set

* diffusers is no longer imported by vllm_omni.diffusion.data (annotation only);
  vllm_omni.diffusion.io_support (diffusion registry + distributed state) is
  imported lazily by config_factory and async_omni_engine, so constructing a
  Qwen3-TTS pipeline config / AsyncOmni no longer loads diffusers,
  x_transformers or whisper (guarded by tests/config/test_import_guard.py).
* requirements/common.txt = core.txt (AR/speech runtime) + media.txt
  (diffusion/video/extra audio); VLLM_OMNI_EDGE_BUILD=1 makes setup.py use
  requirements/<device>-edge.txt (core only); `edge` extra adds websockets and
  sounddevice for on-device clients.

Measured on this host (CUDA venv): `import vllm` 7.0-7.9 s, `import vllm_omni`
adds ~20 s of which ~17 s is transformers' own import (transformers.utils.import_utils)
pulled in by vllm_omni.patch / vllm_omni.version; that floor is not
vLLM-Omni's and is unchanged. Guard test: no diffusion module loaded on the TTS path.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

pyproject.toml

@@ -35,6 +35,12 @@ classifiers = [
 
 [project.optional-dependencies]
 
+# Edge / on-device clients (codec-token streaming client, playback).
+edge = [
+    "websockets>=13.0",
+    "sounddevice>=0.5.0",
+]
+
 dev = [
     "pytest==9.1.1",
     "pytest-asyncio==1.4.0",

requirements/common.txt

@@ -1,41 +1,3 @@
 # Common dependencies for all platforms
-# transformers: required unified multimodal model classes are unavailable
-# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
-# regression. transformers 5.9 also changed
-# _adjust_missing_and_unexpected_keys from
-# list+list to set|set, breaking models that declare
-# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
-# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
-# That breakage is now handled in code instead of by pinning:
-# MossTTSNanoForGeneration wraps weight loading in
-# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 5.10.1, < 5.15
-av>=14.0.0
-omegaconf>=2.3.0
-diffusers==0.40.0
-# Hub kernels, including LTX-2.5 NATTEN. Keep this pin in sync with the
-# supported transformers range above.
-kernels==0.15.2
-safetensors>=0.8.0
-accelerate==1.12.0
-soundfile>=0.13.1
-cache-dit==1.5.0
-tqdm>=4.66.0
-torchsde>=0.2.6
-openai-whisper>=20250625
-imageio[ffmpeg]>=2.37.2
-x-transformers>=2.12.2
-einops>=0.8.1
-prettytable>=3.8.0
-aenum==3.1.16
-pyzmq>=25.0.0
-janus>=1.0.0
-msgpack>=1.0.0
-pydantic>=2.1.0
-pydantic-settings>=2.0.0
-# Required for GGUF-quantized diffusion model loading (gguf_adapters). The CI
-# base image does not ship it, so declare it explicitly. The loader also guards
-# this import so a missing gguf no longer breaks all diffusion model loading.
-gguf>=0.10.0
-# Required for running the Cosmos3 model with guardrails enabled (default behavior).
-cosmos-guardrail>=0.3.1
+# -r core.txt
+-r media.txt

requirements/core.txt

@@ -0,0 +1,17 @@
+# Core runtime deps for AR / generation (TTS, speech) pipelines. Media /
+# diffusion / video deps live in media.txt; common.txt pulls both.
+transformers >= 5.10.1, < 5.15
+av>=14.0.0
+omegaconf>=2.3.0
+safetensors>=0.8.0
+accelerate==1.12.0
+soundfile>=0.13.1
+tqdm>=4.66.0
+einops>=0.8.1
+prettytable>=3.8.0
+aenum==3.1.16
+pyzmq>=25.0.0
+janus>=1.0.0
+msgpack>=1.0.0
+pydantic>=2.1.0
+pydantic-settings>=2.0.0

requirements/cpu-edge.txt

@@ -0,0 +1,4 @@
+# AR-only (TTS/speech) CPU build: VLLM_OMNI_EDGE_BUILD=1 VLLM_TARGET_DEVICE=cpu pip install -e .
+-r core.txt
+psutil
+intel-openmp; platform_machine == "x86_64"

requirements/cuda-edge.txt

@@ -0,0 +1,2 @@
+# AR-only (TTS/speech) CUDA build: VLLM_OMNI_EDGE_BUILD=1 pip install -e .
+-r core.txt

requirements/media.txt

@@ -0,0 +1,10 @@
+# Diffusion / image / video / extra-audio deps (not needed by AR-only edge builds).
+diffusers==0.40.0
+kernels==0.15.2
+cache-dit==1.5.0
+torchsde>=0.2.6
+openai-whisper>=20250625
+imageio[ffmpeg]>=2.37.2
+x-transformers>=2.12.2
+gguf>=0.10.0
+cosmos-guardrail>=0.3.1

setup.py

@@ -227,6 +227,12 @@ def get_install_requires() -> list[str]:
     device = detect_target_device()
     requirements_dir = Path(__file__).parent / "requirements"
     requirements_file = requirements_dir / f"{device}.txt"
+    # VLLM_OMNI_EDGE_BUILD=1 selects the AR-only dependency set (no diffusion /
+    # video stack) when requirements/<device>-edge.txt exists.
+    if os.environ.get("VLLM_OMNI_EDGE_BUILD", "0") == "1":
+        edge_file = requirements_dir / f"{device}-edge.txt"
+        if edge_file.exists():
+            requirements_file = edge_file
 
     print(f"Loading requirements from: {requirements_file}")
     requirements = load_requirements(requirements_file)

tests/config/test_import_guard.py

@@ -0,0 +1,36 @@
+"""Import guard for AR/TTS-only deployments: heavy diffusion dependencies stay unloaded (CPU).
+
+Runs in a subprocess so the assertion is not polluted by modules other tests
+imported into this interpreter.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_PROBE = r"""
+import json, sys, time
+t = time.perf_counter(); import vllm; t_vllm = time.perf_counter() - t
+t = time.perf_counter(); import vllm_omni; t_omni = time.perf_counter() - t
+from vllm_omni.config.config_factory import StageConfigFactory  # noqa: F401
+from vllm_omni.entrypoints.async_omni import AsyncOmni  # noqa: F401
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+assert "qwen3_tts" in OMNI_PIPELINES
+heavy = ["diffusers", "x_transformers", "whisper", "vllm_omni.diffusion.registry", "vllm_omni.diffusion.io_support"]
+print(json.dumps({"t_vllm": t_vllm, "t_omni": t_omni, "loaded": [m for m in heavy if m in sys.modules]}))
+"""
+
+
+def test_tts_path_does_not_import_diffusion_stack():
+    out = subprocess.run([sys.executable, "-c", _PROBE], capture_output=True, text=True, timeout=600)
+    assert out.returncode == 0, out.stderr[-2000:]
+    info = json.loads(out.stdout.strip().splitlines()[-1])
+    assert info["loaded"] == [], f"heavy modules imported on the AR/TTS path: {info['loaded']}"
+    # Import overhead of vllm_omni over bare vllm is reported (not gated: it is host dependent).
+    print(f"import vllm {info['t_vllm']:.1f}s, +vllm_omni {info['t_omni']:.1f}s")

vllm_omni/config/config_factory.py

@@ -32,12 +32,19 @@ from vllm_omni.config.stage_config import (
     normalize_pipeline_cli_overrides,
 )
 from vllm_omni.config.yaml_util import create_config
-from vllm_omni.diffusion.io_support import get_diffusion_output_type
 from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
 logger = init_logger(__name__)
 
 
+def _get_diffusion_output_type(model_class_name: str):
+    # Lazy: vllm_omni.diffusion.io_support pulls in the diffusion registry and
+    # distributed state; AR-only (e.g. TTS) deployments never need it.
+    from vllm_omni.diffusion.io_support import get_diffusion_output_type
+
+    return get_diffusion_output_type(model_class_name)
+
+
 # Default degree for any parallel axis / replica count that isn't set anywhere
 # (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank.
 # TODO(composable_parallel): this "1" default and the per-axis device-layout
@@ -642,7 +649,7 @@ class StageConfigFactory:
             },
             "engine_args": create_config(engine_args),
             "final_output": True,
-            "final_output_type": get_diffusion_output_type(model_class_name),
+            "final_output_type": _get_diffusion_output_type(model_class_name),
         }
 
         return [config_dict]

vllm_omni/diffusion/data.py

@@ -10,7 +10,9 @@ from dataclasses import dataclass, field, fields
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
-import diffusers
+if TYPE_CHECKING:  # heavy; only needed for annotations here
+    import diffusers
+
 import huggingface_hub
 import torch
 from PIL import Image
@@ -897,7 +899,7 @@ class OmniDiffusionConfig:
     # kwargs forwarded to pipeline.__call__()
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
     # Actual diffusers pipeline object (to determine inputs of the dummy run)
-    diffusers_pipeline_cls: type[diffusers.DiffusionPipeline] | None = None  # pyright: ignore[reportPrivateImportUsage]
+    diffusers_pipeline_cls: "type[diffusers.DiffusionPipeline] | None" = None  # pyright: ignore[reportPrivateImportUsage]
 
     # http server endpoint config, would be ignored in local mode
     host: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -45,7 +45,6 @@ from vllm_omni.diffusion.data import (
     parse_attention_config,
     resolve_model_class_name,
 )
-from vllm_omni.diffusion.io_support import get_diffusion_output_type
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_engine_utils import (
     SHUTDOWN_ENQUEUE_TIMEOUT_S,
@@ -1141,6 +1140,8 @@ class AsyncOmniEngine:
         num_devices = max(1, int(parallel_config.world_size))
         devices = ",".join(str(i) for i in range(num_devices))
         model_class_name = kwargs.get("model_class_name", None)
+        from vllm_omni.diffusion.io_support import get_diffusion_output_type
+
         final_output_type = get_diffusion_output_type(model_class_name)
 
         attention_config = None