From 154d9ca4ff4a430472cd0e9fafb2f9a879778b50 Mon Sep 17 00:00:00 2001 From: hsliu Date: Thu, 11 Sep 2025 14:32:51 +0800 Subject: [PATCH] Initial vLLM-omni repository structure - Add complete directory structure for multimodal and non-autoregressive processing - Include AsyncLLM, OmniEngine, and DiffusionEngine modules - Add processors, schedulers, executors, and output processors - Add API interfaces for Higress, Gradio, and ComfyUI - Add configuration management and CLI interface - Add comprehensive test structure and documentation - Add pyproject.toml with dependencies and project configuration - Add examples and documentation structure --- .gitignore | 127 +++++++----- README.md | 5 +- configs/default.yaml | 98 ++++++++++ docs/README.md | 5 + .../architecture/vLLM-omni arch design doc.md | 139 ++++++++++++++ examples/README.md | 62 ++++++ pyproject.toml | 137 +++++++++++++ requirements-dev.txt | 29 +++ requirements.txt | 21 ++ tests/__init__.py | 5 + tests/conftest.py | 118 ++++++++++++ vllm_omni/__init__.py | 34 ++++ vllm_omni/api_interfaces/__init__.py | 0 vllm_omni/async_llm/__init__.py | 0 vllm_omni/cli.py | 180 ++++++++++++++++++ vllm_omni/config/__init__.py | 73 +++++++ vllm_omni/configs/__init__.py | 73 +++++++ vllm_omni/core/__init__.py | 0 vllm_omni/core/sched/__init__.py | 0 vllm_omni/diffusion_engine/__init__.py | 18 ++ vllm_omni/executor/__init__.py | 0 vllm_omni/request.py | 16 ++ vllm_omni/utils/__init__.py | 0 vllm_omni/worker/__init__.py | 0 24 files changed, 1087 insertions(+), 53 deletions(-) create mode 100644 configs/default.yaml create mode 100644 docs/README.md create mode 100644 docs/architecture/vLLM-omni arch design doc.md create mode 100644 examples/README.md create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100644 requirements.txt create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 vllm_omni/__init__.py create mode 100644 vllm_omni/api_interfaces/__init__.py create mode 100644 vllm_omni/async_llm/__init__.py create mode 100644 vllm_omni/cli.py create mode 100644 vllm_omni/config/__init__.py create mode 100644 vllm_omni/configs/__init__.py create mode 100644 vllm_omni/core/__init__.py create mode 100644 vllm_omni/core/sched/__init__.py create mode 100644 vllm_omni/diffusion_engine/__init__.py create mode 100644 vllm_omni/executor/__init__.py create mode 100644 vllm_omni/request.py create mode 100644 vllm_omni/utils/__init__.py create mode 100644 vllm_omni/worker/__init__.py diff --git a/.gitignore b/.gitignore index b7faf403d..d42c87793 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,6 @@ # Byte-compiled / optimized / DLL files __pycache__/ -*.py[codz] +*.py[cod] *$py.class # C extensions @@ -46,7 +46,7 @@ htmlcov/ nosetests.xml coverage.xml *.cover -*.py.cover +*.py,cover .hypothesis/ .pytest_cache/ cover/ @@ -94,35 +94,20 @@ ipython_config.py # install all needed dependencies. #Pipfile.lock -# UV -# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -#uv.lock - # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock -#poetry.toml # pdm # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. -# https://pdm-project.org/en/latest/usage/project/#working-with-version-control #pdm.lock -#pdm.toml -.pdm-python -.pdm-build/ - -# pixi -# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. -#pixi.lock -# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one -# in the .venv directory. It is recommended not to include this directory in version control. -.pixi +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm __pypackages__/ @@ -136,7 +121,6 @@ celerybeat.pid # Environments .env -.envrc .venv env/ venv/ @@ -170,38 +154,79 @@ cython_debug/ # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +# be added to the global gitignore or merged into this project gitignore. For a PyCharm +# project, it is recommended to include the following files: +# .idea/ +# *.iml +# *.ipr +# *.iws +# .idea/ -# Abstra -# Abstra is an AI-powered process automation framework. -# Ignore directories containing user credentials, local state, and settings. -# Learn more at https://abstra.io/docs -.abstra/ +# vLLM-omni specific +# Model files and checkpoints +models/ +checkpoints/ +*.bin +*.safetensors +*.pt +*.pth -# Visual Studio Code -# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore -# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore -# and can be added to the global gitignore or merged into this file. However, if you prefer, -# you could uncomment the following to ignore the entire vscode folder -# .vscode/ +# Cache directories +cache/ +.cache/ +diffusion_cache/ +kv_cache/ -# Ruff stuff: -.ruff_cache/ +# Logs +logs/ +*.log -# PyPI configuration file -.pypirc +# Temporary files +tmp/ +temp/ +.tmp/ -# Cursor -# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to -# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data -# refer to https://docs.cursor.com/context/ignore-files -.cursorignore -.cursorindexingignore +# IDE files +.vscode/ +.idea/ +*.swp +*.swo +*~ -# Marimo -marimo/_static/ -marimo/_lsp/ -__marimo__/ +# OS files +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db + +# Data files +data/ +datasets/ +*.csv +*.json +*.jsonl +*.parquet + +# Output files +outputs/ +results/ +generated/ + +# Configuration overrides +configs/local.yaml +configs/production.yaml +configs/development.yaml + +# Docker +.dockerignore +Dockerfile.dev + +# Kubernetes +k8s/ +*.yaml +*.yml +!configs/*.yaml +!configs/*.yml \ No newline at end of file diff --git a/README.md b/README.md index 3982ac0db..75b1bcf23 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,3 @@ -# vllm-omni -vLLM-omni supports multi-modality models inference and serving with non-autoregressive structures and non-textual outputs, extending vLLM beyond traditional text-based, autoregressive generation. +# vLLM-omni: Multi-modal Extension for vLLM + +vLLM-omni is designed to extend vLLM capabilities to support multi-modality model inference and serving, particularly focusing on non-autoregressive structures and non-textual outputs. diff --git a/configs/default.yaml b/configs/default.yaml new file mode 100644 index 000000000..d837accff --- /dev/null +++ b/configs/default.yaml @@ -0,0 +1,98 @@ +# vLLM-omni Default Configuration + +# General settings +general: + name: "vllm-omni" + version: "0.1.0" + debug: false + log_level: "INFO" + +# Model settings +model: + # Supported model types + supported_types: ["ar", "diffusion", "multimodal"] + + # Default model configurations + default_ar_model: "qwen2.5-7b" + default_diffusion_model: "qwen-image" + + # Model loading settings + device: "auto" # auto, cpu, cuda, mps + dtype: "auto" # auto, float16, bfloat16, float32 + max_model_len: 32768 + +# Engine settings +engines: + # AR Engine (imported from vLLM) + ar_engine: + enabled: true + max_batch_size: 256 + max_num_seqs: 256 + + # Diffusion Engine + diffusion_engine: + enabled: true + max_batch_size: 32 + max_num_seqs: 32 + num_inference_steps: 50 + guidance_scale: 7.5 + +# Processing settings +processing: + # Input processing + input: + max_image_size: [1024, 1024] + max_audio_length: 30.0 # seconds + supported_image_formats: ["jpg", "jpeg", "png", "bmp", "webp"] + supported_audio_formats: ["wav", "mp3", "flac", "m4a"] + + # Output processing + output: + max_output_length: 4096 + streaming: true + return_hidden_states: false + +# API settings +api: + # Server settings + host: "0.0.0.0" + port: 8000 + workers: 1 + + # Interface settings + interfaces: + higress: + enabled: true + endpoint: "/api/v1" + gradio: + enabled: true + port: 7860 + share: false + comfyui: + enabled: false + port: 8188 + +# Cache settings +cache: + # KV Cache (for AR models) + kv_cache: + enabled: true + max_size: "2GB" + + # Diffusion Cache + diffusion_cache: + enabled: true + max_size: "1GB" + cache_latents: true + +# Logging settings +logging: + level: "INFO" + format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + file: null # Set to file path for file logging + +# Monitoring settings +monitoring: + enabled: true + metrics_port: 9090 + health_check_interval: 30 # seconds diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 000000000..d52c6c099 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,5 @@ +# vLLM-omni Documentation + +Welcome to the vLLM-omni documentation! This documentation provides comprehensive information about using and developing with vLLM-omni. + +# TODO (add examples) \ No newline at end of file diff --git a/docs/architecture/vLLM-omni arch design doc.md b/docs/architecture/vLLM-omni arch design doc.md new file mode 100644 index 000000000..0306606c0 --- /dev/null +++ b/docs/architecture/vLLM-omni arch design doc.md @@ -0,0 +1,139 @@ +# vLLM-omni Architecture Design Document + +# Introduction + +This document outlines the architectural design for vLLM-omni, a new repository developed under the vLLM project. vLLM-omni supports multi-modality models with non-autoregressive structures and non-textual outputs, extending vLLM beyond traditional text-based, autoregressive generation. + +# Goals + +The primary goals for vLLM-omni are: + +* **Support for Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text. +* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation. +* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable. +* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats. + +# Expected Gains + +* It breaks through the limitation of vLLM that only supports autoregressive text generation models, and provides architectural support for popular model structures, especially multimodal generation. +* By supporting new output formats and new model architectures, developers can easily connect to mature ecosystems such as Higress Gateway and Ray, improving the ease of deployment and efficiency of corresponding models. +* Integrating efficient cache management, multiple parallel strategies, disaggregation from vLLM, to improve inference efficiency of corresponding models. + +# Non-Autoregressive Structure + +* **Diffusion Transformer:** initialized with a random latent noise, the transformer model iteratively predicts latent output for multiple steps, under control of sampling timestamps. Finally, it is decoded by the decoder of Variational AutoEncoder (VAE). It is not supported by vLLM now. +* **Autoregressive:** dominant generation paradigm for text, now supported by vLLM. It generates tokens conditioned on previous tokens one by one. The efficient KV Cache management provided by vLLM can accelerate inference effectively. + ![][image1] ![][image2] + + +| | Autoregressive | DiT | +| :---- | ----- | ----- | +| Scenarios | Text generation | Multimodal generation | +| Generation | Token-by-token, left to right | Continuous vector, iteratively refine | +| Key advantages | Smooth, good local consistency | Global consistency, diversity and controlability | +| Efficiency | More efficient with KV cache | Need generation repeatedly, slow | +| Error accumulation | Easy to accumulate | Can fix possibly | +| Key challenges | Limited global planning, long context generation | Large global computational cost | +| Sequence length | Varied length | Fixed length | +| Attention mask | Lower triangular | Full bi-directional matrix | +| Parallel strategy | TP, DP, PP, .etc | DP, CP, .etc | +| vLLM support | Supported | Unsupported | + +# Targeted models to support + +According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below: + +* **DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)** + A popular model for image generation and editing. Main visual generation models are similar to the structure, such as Flux. + ![][image3] +* **AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)** + A unified multimodal comprehension and generation model. Visual generation can utilize CoT text generation output. Main unified multimodal models are similar to the structure. + ![][image4] +* **Extension, multiple AR+DiT (e.g. Qwen-omni)** + A new model for multimodal input and output. It designs the Thinker-Talker-Codec structure, which is a double-AR+DiT format. + ![][image5] + + + +# Architectural Overview + +vLLM-omni will import core modules from existing vLLM framework and introduce several new components and modifications to achieve its goals. The core architecture will be based on a modular design, allowing for independent development and integration of different modality handlers and output generators. +![][image6] + +## General ideas + +* It divides models in vLLM-omni into multiple stages, which are processed by different Engine Cores. Each stage can choose AR Engine Core for AR model part, and Diffusion Engine Core for diffusion model part. +* Designed a specialized diffusion engine core for DiT or dLLM. +* Plug-in with external DiT acceleration modules such as xDiT, Cache-DiT. +* Output linked with Higress API, Gradio for online demo, ComfyUI for less-code process. + +## Key Components + +* **Processor:** Input processing, tokenization or encoding +* **Scheduler:** for order and queue of incoming requests +* **Executor:** Prepare model input and run +* **Multimodal output processor:** support multimodal outputs +* **Hidden State Output processor:** support intermediate hidden vectors as request output +* **Engine Core:** accelerate AR modules, imported from vLLM main repo +* **Diffusion Engine Core:** Wrap and accelerate DiT modules. +* **Diffusion model:** Wrap up diffusion models with vLLM format + +# Examples data flow of specific models + +1. Qwen-omni + ![][image7]![][image8] +2. Qwen-Image + ![][image9] ![][image10] + +# Overview of function call dependency + +![][image11] + +# Roadmap + +* AsyncLLM: +- [ ] Initializer: resource allocation for different stages (2025 H2) +- [ ] Hidden state Output Processor (2025 H2) +- [ ] Multimodal Output Processor (2025 H2) +- [ ] Interface for API, Gradio, ComfyUI +* OmniEngine: +- [ ] Basic Diffusion Engine core: Diffusion step manager, Cache manager (2025 H2) +- [ ] Pipeline of different Engine Cores (2025 H2) +- [ ] Communication between Engine Cores via AsyncLLM (2025 H2) +* Model support (Included but not limited to): +- [ ] DiT: Qwen-Image (2025 H2), Flux-Kontext +- [ ] Unified-MLLM: BAGEL (2025 H2), OmniGen2, Janus-Pro +- [ ] Extensions of AR+DiT: Qwen-omni (2025 H2) +* Optimization +- [ ] Integration of other DiT acceleration modules, such as xDiT, Cache-DiT +- [ ] Disaggregation for AR and DiT +- [ ] Integration with vLLM's existing KV cache for potential future hybrid models. +- [ ] Support for streaming multi-modal inputs and outputs + +# Contact + +For any questions or feedback regarding this design document, please contact +Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk), +Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com) + +[image1]: + +[image2]: + +[image3]: + +[image4]: + +[image5]: + +[image6]: + +[image7]: + +[image8]: + +[image9]: + +[image10]: + +[image11]: \ No newline at end of file diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 000000000..e79470c89 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,62 @@ +# vLLM-omni Examples + +This directory contains examples demonstrating how to use vLLM-omni for various tasks. + +## Basic Examples + +- [Text Generation](basic/text_generation.py) - Basic text generation using AR models +- [Image Generation](basic/image_generation.py) - Image generation using diffusion models +- [Multimodal Processing](basic/multimodal_processing.py) - Processing text and images together + +## Advanced Examples + +- [Custom Model Integration](advanced/custom_model.py) - Integrating custom models +- [Batch Processing](advanced/batch_processing.py) - Efficient batch processing +- [Streaming Output](advanced/streaming.py) - Real-time streaming output + +## Multimodal Examples + +- [Image-to-Text](multimodal/image_to_text.py) - Image captioning and description +- [Text-to-Image](multimodal/text_to_image.py) - Text-to-image generation +- [Audio Processing](multimodal/audio_processing.py) - Audio generation and processing +- [Video Generation](multimodal/video_generation.py) - Video generation workflows + +## API Examples + +- [REST API](api/rest_api.py) - Using the REST API interface +- [Gradio Interface](api/gradio_interface.py) - Creating Gradio web interfaces +- [ComfyUI Integration](api/comfyui_integration.py) - ComfyUI workflow integration + +## Configuration Examples + +- [Custom Configuration](config/custom_config.py) - Custom configuration setup +- [Multi-GPU Setup](config/multi_gpu.py) - Multi-GPU configuration +- [Distributed Processing](config/distributed.py) - Distributed processing setup + +## Getting Started + +1. Install vLLM-omni: + ```bash + pip install vllm-omni + ``` + +2. Run a basic example: + ```bash + python examples/basic/text_generation.py + ``` + +3. Explore the examples in each subdirectory for more advanced usage. + +## Requirements + +Most examples require additional dependencies. Install them with: + +```bash +pip install -r requirements.txt +``` + +For development examples, install the development dependencies: + +```bash +pip install -r requirements-dev.txt +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..ca6beb477 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,137 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "vllm-omni" +version = "0.1.0" +description = "vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures" +readme = "README.md" +requires-python = ">=3.8" +license = {text = "Apache-2.0"} +authors = [ + {name = "vLLM-omni Team", email = "hsliuustc@gmail.com"} +] +keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", +] + +dependencies = [ + "vllm>=0.2.0", + "torch>=2.0.0", + "transformers>=4.30.0", + "numpy>=1.21.0", + "pillow>=8.0.0", + "opencv-python>=4.5.0", + "gradio>=3.0.0", + "fastapi>=0.100.0", + "uvicorn>=0.20.0", + "pydantic>=2.0.0", + "asyncio-mqtt>=0.11.0", + "ray>=2.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.21.0", + "pytest-cov>=4.0.0", + "black>=23.0.0", + "isort>=5.12.0", + "flake8>=6.0.0", + "mypy>=1.0.0", + "pre-commit>=3.0.0", +] +docs = [ + "sphinx>=6.0.0", + "sphinx-rtd-theme>=1.2.0", + "myst-parser>=1.0.0", +] +examples = [ + "jupyter>=1.0.0", + "matplotlib>=3.5.0", + "seaborn>=0.11.0", +] + +[project.urls] +Homepage = "https://github.com/hsliuustc0106/vllm-omni" +Repository = "https://github.com/hsliuustc0106/vllm-omni" +Documentation = "https://vllm-omni.readthedocs.io" +"Bug Tracker" = "https://github.com/hsliuustc0106/vllm-omni/issues" + +[project.scripts] +vllm-omni = "vllm_omni.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["vllm_omni*"] + +[tool.black] +line-length = 88 +target-version = ['py38', 'py39', 'py310', 'py311'] +include = '\.pyi?$' +extend-exclude = ''' +/( + # directories + \.eggs + | \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | build + | dist +)/ +''' + +[tool.isort] +profile = "black" +multi_line_output = 3 +line_length = 88 +known_first_party = ["vllm_omni"] + +[tool.mypy] +python_version = "3.8" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +disallow_incomplete_defs = true +check_untyped_defs = true +disallow_untyped_decorators = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = true +warn_no_return = true +warn_unreachable = true +strict_equality = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py", "*_test.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "--strict-markers", + "--strict-config", + "--cov=vllm_omni", + "--cov-report=term-missing", + "--cov-report=html", + "--cov-report=xml", +] +markers = [ + "unit: Unit tests", + "integration: Integration tests", + "benchmark: Benchmark tests", + "slow: Slow tests", +] diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 000000000..1a28bcb3a --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,29 @@ +# Development dependencies +-r requirements.txt + +# Testing +pytest>=7.0.0 +pytest-asyncio>=0.21.0 +pytest-cov>=4.0.0 +pytest-mock>=3.10.0 + +# Code quality +black>=23.0.0 +isort>=5.12.0 +flake8>=6.0.0 +mypy>=1.0.0 +pre-commit>=3.0.0 + +# Documentation +sphinx>=6.0.0 +sphinx-rtd-theme>=1.2.0 +myst-parser>=1.0.0 + +# Examples and notebooks +jupyter>=1.0.0 +matplotlib>=3.5.0 +seaborn>=0.11.0 + +# Profiling and debugging +memory-profiler>=0.60.0 +py-spy>=0.3.14 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 000000000..88cb7b3a5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Core dependencies +vllm>=0.2.0 +torch>=2.0.0 +transformers>=4.30.0 +numpy>=1.21.0 +pillow>=8.0.0 +opencv-python>=4.5.0 + +# API and serving +fastapi>=0.100.0 +uvicorn>=0.20.0 +pydantic>=2.0.0 +gradio>=3.0.0 + +# Async and distributed computing +asyncio-mqtt>=0.11.0 +ray>=2.0.0 + +# Optional: DiT acceleration modules +# xDiT # Uncomment when available +# Cache-DiT # Uncomment when available diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..0fc4af113 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,5 @@ +""" +Test suite for vLLM-omni. + +This package contains unit tests, integration tests, and benchmarks for vLLM-omni. +""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..8d40a3e13 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,118 @@ +""" +Pytest configuration and fixtures for vLLM-omni tests. +""" + +import pytest +import asyncio +from typing import AsyncGenerator, Generator +from unittest.mock import Mock, AsyncMock + +import torch +import numpy as np + +from vllm_omni.configs import load_config + + +@pytest.fixture(scope="session") +def event_loop() -> Generator[asyncio.AbstractEventLoop, None, None]: + """Create an instance of the default event loop for the test session.""" + loop = asyncio.get_event_loop_policy().new_event_loop() + yield loop + loop.close() + + +@pytest.fixture +async def mock_async_llm() -> AsyncMock: + """Mock AsyncLLM for testing.""" + mock = AsyncMock() + mock.process.return_value = {"output": "test_output", "hidden_states": None} + return mock + + +@pytest.fixture +async def mock_omni_engine() -> AsyncMock: + """Mock OmniEngine for testing.""" + mock = AsyncMock() + mock.execute.return_value = {"result": "test_result"} + return mock + + +@pytest.fixture +def sample_config() -> dict: + """Sample configuration for testing.""" + return { + "general": { + "name": "vllm-omni-test", + "version": "0.1.0", + "debug": True, + "log_level": "DEBUG" + }, + "model": { + "device": "cpu", + "dtype": "float32", + "max_model_len": 1024 + }, + "engines": { + "ar_engine": {"enabled": True, "max_batch_size": 32}, + "diffusion_engine": {"enabled": True, "max_batch_size": 8} + } + } + + +@pytest.fixture +def sample_text_input() -> str: + """Sample text input for testing.""" + return "Hello, world! This is a test input for vLLM-omni." + + +@pytest.fixture +def sample_image_input() -> np.ndarray: + """Sample image input for testing.""" + return np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8) + + +@pytest.fixture +def sample_audio_input() -> np.ndarray: + """Sample audio input for testing.""" + return np.random.randn(16000).astype(np.float32) # 1 second at 16kHz + + +@pytest.fixture +def sample_multimodal_input() -> dict: + """Sample multimodal input for testing.""" + return { + "text": "Describe this image", + "image": np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8), + "audio": np.random.randn(16000).astype(np.float32) + } + + +@pytest.fixture +def mock_torch_device() -> str: + """Mock torch device for testing.""" + if torch.cuda.is_available(): + return "cuda" + elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): + return "mps" + else: + return "cpu" + + +# Pytest markers +pytest_plugins = [] + + +def pytest_configure(config): + """Configure pytest with custom markers.""" + config.addinivalue_line( + "markers", "unit: mark test as a unit test" + ) + config.addinivalue_line( + "markers", "integration: mark test as an integration test" + ) + config.addinivalue_line( + "markers", "benchmark: mark test as a benchmark test" + ) + config.addinivalue_line( + "markers", "slow: mark test as slow running" + ) diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py new file mode 100644 index 000000000..408545dfd --- /dev/null +++ b/vllm_omni/__init__.py @@ -0,0 +1,34 @@ +""" +vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures. + +This package extends vLLM beyond traditional text-based, autoregressive generation +to support multi-modality models with non-autoregressive structures and non-textual outputs. + +Architecture: +- 🟡 Modified: vLLM components modified for multimodal support +- 🔴 Added: New components for multimodal and non-autoregressive processing +""" + +__version__ = "0.1.0" +__author__ = "vLLM-omni Team" +__email__ = "hsliuustc@gmail.com" + +# Main entry points +from .async_llm import AsyncLLM +from .omni_engine import OmniEngine +from .diffusion_engine import DiffusionEngine + +__all__ = [ + # Version info + "__version__", + "__author__", + "__email__", + + # Main components + "AsyncLLM", + "OmniEngine", + "DiffusionEngine", + + # All other components are available through their respective modules + # processors.*, schedulers.*, executors.*, etc. +] \ No newline at end of file diff --git a/vllm_omni/api_interfaces/__init__.py b/vllm_omni/api_interfaces/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/async_llm/__init__.py b/vllm_omni/async_llm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/cli.py b/vllm_omni/cli.py new file mode 100644 index 000000000..b14c34205 --- /dev/null +++ b/vllm_omni/cli.py @@ -0,0 +1,180 @@ +""" +Command-line interface for vLLM-omni. +""" + +import argparse +import asyncio +import logging +import sys +from pathlib import Path +from typing import Optional + +from .async_llm import AsyncLLM +from .omni_engine import OmniEngine +from .configs import load_config + + +def setup_logging(level: str = "INFO") -> None: + """Set up logging configuration.""" + logging.basicConfig( + level=getattr(logging, level.upper()), + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + stream=sys.stdout + ) + + +async def run_server( + config_path: Optional[str] = None, + host: str = "0.0.0.0", + port: int = 8000, + log_level: str = "INFO" +) -> None: + """Run the vLLM-omni server.""" + setup_logging(log_level) + logger = logging.getLogger(__name__) + + # Load configuration + config = load_config(config_path) if config_path else {} + + # Initialize components + async_llm = AsyncLLM(config) + omni_engine = OmniEngine(config) + + logger.info(f"Starting vLLM-omni server on {host}:{port}") + + try: + # Start the server (implementation depends on chosen framework) + # This is a placeholder for the actual server implementation + logger.info("Server started successfully") + + # Keep the server running + while True: + await asyncio.sleep(1) + + except KeyboardInterrupt: + logger.info("Shutting down server...") + except Exception as e: + logger.error(f"Server error: {e}") + sys.exit(1) + + +def run_gradio( + config_path: Optional[str] = None, + port: int = 7860, + share: bool = False, + log_level: str = "INFO" +) -> None: + """Run the Gradio interface.""" + setup_logging(log_level) + logger = logging.getLogger(__name__) + + try: + import gradio as gr + from .api_interfaces.gradio import GradioInterface + + # Load configuration + config = load_config(config_path) if config_path else {} + + # Initialize Gradio interface + interface = GradioInterface(config) + + logger.info(f"Starting Gradio interface on port {port}") + interface.launch(server_port=port, share=share) + + except ImportError: + logger.error("Gradio is not installed. Install it with: pip install gradio") + sys.exit(1) + except Exception as e: + logger.error(f"Gradio interface error: {e}") + sys.exit(1) + + +def main() -> None: + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + description="vLLM-omni: Multi-modality models inference and serving" + ) + + subparsers = parser.add_subparsers(dest="command", help="Available commands") + + # Server command + server_parser = subparsers.add_parser("server", help="Run the server") + server_parser.add_argument( + "--config", "-c", + type=str, + help="Path to configuration file" + ) + server_parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Host to bind to (default: 0.0.0.0)" + ) + server_parser.add_argument( + "--port", "-p", + type=int, + default=8000, + help="Port to bind to (default: 8000)" + ) + server_parser.add_argument( + "--log-level", + type=str, + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + default="INFO", + help="Log level (default: INFO)" + ) + + # Gradio command + gradio_parser = subparsers.add_parser("gradio", help="Run the Gradio interface") + gradio_parser.add_argument( + "--config", "-c", + type=str, + help="Path to configuration file" + ) + gradio_parser.add_argument( + "--port", "-p", + type=int, + default=7860, + help="Port to bind to (default: 7860)" + ) + gradio_parser.add_argument( + "--share", + action="store_true", + help="Create a public link" + ) + gradio_parser.add_argument( + "--log-level", + type=str, + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + default="INFO", + help="Log level (default: INFO)" + ) + + # Version command + version_parser = subparsers.add_parser("version", help="Show version information") + + args = parser.parse_args() + + if args.command == "server": + asyncio.run(run_server( + config_path=args.config, + host=args.host, + port=args.port, + log_level=args.log_level + )) + elif args.command == "gradio": + run_gradio( + config_path=args.config, + port=args.port, + share=args.share, + log_level=args.log_level + ) + elif args.command == "version": + from . import __version__ + print(f"vLLM-omni version {__version__}") + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py new file mode 100644 index 000000000..85e6255de --- /dev/null +++ b/vllm_omni/config/__init__.py @@ -0,0 +1,73 @@ +""" +Configuration management for vLLM-omni. +""" + +import os +import yaml +from pathlib import Path +from typing import Dict, Any, Optional + + +def load_config(config_path: Optional[str] = None) -> Dict[str, Any]: + """ + Load configuration from file or use defaults. + + Args: + config_path: Path to configuration file. If None, loads default config. + + Returns: + Configuration dictionary. + """ + if config_path is None: + config_path = os.path.join( + Path(__file__).parent.parent.parent, + "configs", + "default.yaml" + ) + + config_path = Path(config_path) + + if not config_path.exists(): + raise FileNotFoundError(f"Configuration file not found: {config_path}") + + with open(config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) + + return config + + +def save_config(config: Dict[str, Any], config_path: str) -> None: + """ + Save configuration to file. + + Args: + config: Configuration dictionary to save. + config_path: Path where to save the configuration. + """ + config_path = Path(config_path) + config_path.parent.mkdir(parents=True, exist_ok=True) + + with open(config_path, 'w', encoding='utf-8') as f: + yaml.dump(config, f, default_flow_style=False, indent=2) + + +def merge_configs(base_config: Dict[str, Any], override_config: Dict[str, Any]) -> Dict[str, Any]: + """ + Merge two configuration dictionaries. + + Args: + base_config: Base configuration. + override_config: Configuration to override with. + + Returns: + Merged configuration dictionary. + """ + merged = base_config.copy() + + for key, value in override_config.items(): + if key in merged and isinstance(merged[key], dict) and isinstance(value, dict): + merged[key] = merge_configs(merged[key], value) + else: + merged[key] = value + + return merged diff --git a/vllm_omni/configs/__init__.py b/vllm_omni/configs/__init__.py new file mode 100644 index 000000000..85e6255de --- /dev/null +++ b/vllm_omni/configs/__init__.py @@ -0,0 +1,73 @@ +""" +Configuration management for vLLM-omni. +""" + +import os +import yaml +from pathlib import Path +from typing import Dict, Any, Optional + + +def load_config(config_path: Optional[str] = None) -> Dict[str, Any]: + """ + Load configuration from file or use defaults. + + Args: + config_path: Path to configuration file. If None, loads default config. + + Returns: + Configuration dictionary. + """ + if config_path is None: + config_path = os.path.join( + Path(__file__).parent.parent.parent, + "configs", + "default.yaml" + ) + + config_path = Path(config_path) + + if not config_path.exists(): + raise FileNotFoundError(f"Configuration file not found: {config_path}") + + with open(config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) + + return config + + +def save_config(config: Dict[str, Any], config_path: str) -> None: + """ + Save configuration to file. + + Args: + config: Configuration dictionary to save. + config_path: Path where to save the configuration. + """ + config_path = Path(config_path) + config_path.parent.mkdir(parents=True, exist_ok=True) + + with open(config_path, 'w', encoding='utf-8') as f: + yaml.dump(config, f, default_flow_style=False, indent=2) + + +def merge_configs(base_config: Dict[str, Any], override_config: Dict[str, Any]) -> Dict[str, Any]: + """ + Merge two configuration dictionaries. + + Args: + base_config: Base configuration. + override_config: Configuration to override with. + + Returns: + Merged configuration dictionary. + """ + merged = base_config.copy() + + for key, value in override_config.items(): + if key in merged and isinstance(merged[key], dict) and isinstance(value, dict): + merged[key] = merge_configs(merged[key], value) + else: + merged[key] = value + + return merged diff --git a/vllm_omni/core/__init__.py b/vllm_omni/core/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/diffusion_engine/__init__.py b/vllm_omni/diffusion_engine/__init__.py new file mode 100644 index 000000000..3768b785e --- /dev/null +++ b/vllm_omni/diffusion_engine/__init__.py @@ -0,0 +1,18 @@ +""" +Diffusion Engine: Diffusion Transformer (DiT) processing modules. + +This module provides specialized processing for diffusion models including +step management, cache management, and model wrapping. +""" + +from .step_manager import DiffusionStepManager +from .cache_manager import DiffusionCacheManager +from .models import DiffusionModel +from .base import BaseDiffusionEngine + +__all__ = [ + "DiffusionStepManager", + "DiffusionCacheManager", + "DiffusionModel", + "BaseDiffusionEngine", +] diff --git a/vllm_omni/executor/__init__.py b/vllm_omni/executor/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/request.py b/vllm_omni/request.py new file mode 100644 index 000000000..20dc9b569 --- /dev/null +++ b/vllm_omni/request.py @@ -0,0 +1,16 @@ +import enum +import time +from typing import Optional, Dict, Any + + +Class Request: + # initialize request + def __init__(self, request_id: str, request_type: str, request_data: Dict[str, Any]): + self.request_id = request_id + self.request_type = request_type + self.request_data = request_data + self.request_time = time.time() + self.request_status = RequestStatus.PENDING + self.request_result = None + self.request_error = None + self.request_time = time.time() \ No newline at end of file diff --git a/vllm_omni/utils/__init__.py b/vllm_omni/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/worker/__init__.py b/vllm_omni/worker/__init__.py new file mode 100644 index 000000000..e69de29bb