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/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.txt b/requirements.txt new file mode 100644 index 000000000..45339bd42 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Core dependencies +vllm>=0.9.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/config/__init__.py b/vllm_omni/config/__init__.py new file mode 100644 index 000000000..5174e4a04 --- /dev/null +++ b/vllm_omni/config/__init__.py @@ -0,0 +1,22 @@ +""" +Configuration management for vLLM-omni. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Any, Optional +from vllm_omni.dit_cache_interface import DiTCacheConfig + +from vllm.config import VllmConfig + + +@dataclass +class OmniConfig: + """ + The configuration for vLLM-omni. + """ + + """vllm config""" + vllm_config: VllmConfig = field(default_factory=VllmConfig) + """DiT cache config""" + dit_cache_config: DiTCacheConfig = field(default_factory=DiTCacheConfig) # DiT cache config \ No newline at end of file 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/dit_cache_manager.py b/vllm_omni/core/dit_cache_manager.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/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/distributed/__init__.py b/vllm_omni/distributed/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/dit_cache_interface.py b/vllm_omni/dit_cache_interface.py new file mode 100644 index 000000000..260acac5f --- /dev/null +++ b/vllm_omni/dit_cache_interface.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass + + +@dataclass +class DiTCacheTensor: + """ + A class for specifying how the workers should initialize the DiT cache. + """ + size: int # the size of the cache tensor in bytes + + +@dataclass +class DiTCacheConfig: + """ + The DiT cache configuration of a model. + """ + """How should model runner initialize the KV cache tensors for each layer""" + dit_cache_tensors: list[DiTCacheTensor] + """ + The DiT cache groups of the model. + For models with only one type of DiT, there is only one group that + contains all layers. + """ + kv_cache_groups: list[DiTCacheTensor] \ No newline at end of file diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py new file mode 100644 index 000000000..3768b785e --- /dev/null +++ b/vllm_omni/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/engine/output_processor.py b/vllm_omni/engine/output_processor.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/entrypoints/__init__.py b/vllm_omni/entrypoints/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py new file mode 100644 index 000000000..e69de29bb 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..0919270af --- /dev/null +++ b/vllm_omni/request.py @@ -0,0 +1,7 @@ +import enum +import time +from typing import Optional, Dict, Any +from vllm.v1.request import Request as vLLMRequest + +class OmniRequest(vLLMRequest): + # pass \ No newline at end of file diff --git a/vllm_omni/sample/__init__.py b/vllm_omni/sample/__init__.py new file mode 100644 index 000000000..e69de29bb 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