diff --git a/pyproject.toml b/pyproject.toml index cd7544af8..a846e1b4d 100644 --- a/pyproject.toml +++ b/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", diff --git a/requirements/common.txt b/requirements/common.txt index bb79e0a6a..4b668fdb9 100644 --- a/requirements/common.txt +++ b/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 diff --git a/requirements/core.txt b/requirements/core.txt new file mode 100644 index 000000000..20a3406ea --- /dev/null +++ b/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 diff --git a/requirements/cpu-edge.txt b/requirements/cpu-edge.txt new file mode 100644 index 000000000..e477281cf --- /dev/null +++ b/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" diff --git a/requirements/cuda-edge.txt b/requirements/cuda-edge.txt new file mode 100644 index 000000000..b5b1c393a --- /dev/null +++ b/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 diff --git a/requirements/media.txt b/requirements/media.txt new file mode 100644 index 000000000..595cdaeec --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index 403d8a416..8ac671f03 100644 --- a/setup.py +++ b/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/-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) diff --git a/tests/config/test_import_guard.py b/tests/config/test_import_guard.py new file mode 100644 index 000000000..339afcdd3 --- /dev/null +++ b/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") diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index b2eb64bca..5ee2024ab 100644 --- a/vllm_omni/config/config_factory.py +++ b/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] diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index fd6f0237d..2e6451842 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 97c000c4b..6a8633970 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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