From a04e9451660245cc55489de8f0f31475e07be7cd Mon Sep 17 00:00:00 2001 From: Gaohan123 Date: Fri, 24 Oct 2025 12:16:48 +0800 Subject: [PATCH] add precommit --- .gitattributes | 44 +++++++ .pre-commit-config.yaml | 29 +++++ docs/api/API_DESIGN_TEMPLATE.md | 34 +++--- docs/api/README.md | 4 +- docs/architecture/high level arch design.md | 6 +- .../implementation_architecture.md | 18 +-- examples/basic/api_client.py | 45 +++---- examples/basic/text_generation.py | 87 +++++++------- examples/omni/ar_dit_diffusers.py | 2 +- pyproject.toml | 4 + scripts/README.md | 1 - scripts/test_serving.sh | 30 ++--- tests/__init__.py | 3 +- tests/conftest.py | 32 ++--- tests/unit/test_config.py | 79 +++++-------- vllm_omni/__init__.py | 29 ++--- vllm_omni/config/__init__.py | 8 +- vllm_omni/config/stage_config.py | 43 ++++--- vllm_omni/core/dit_cache_manager.py | 89 +++++++------- vllm_omni/core/sched/scheduler.py | 5 +- vllm_omni/engine/diffusion_engine.py | 19 +-- vllm_omni/entrypoints/api_server.py | 86 +++++++------- vllm_omni/entrypoints/cli/main.py | 4 +- vllm_omni/entrypoints/cli/serve.py | 110 ++++++++---------- 24 files changed, 434 insertions(+), 377 deletions(-) create mode 100644 .gitattributes create mode 100644 .pre-commit-config.yaml diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..f837fd43e --- /dev/null +++ b/.gitattributes @@ -0,0 +1,44 @@ +# Set default behavior to automatically normalize line endings +* text=auto + +# Explicitly declare text files you want to always be normalized and converted +# to native line endings on checkout +*.py text eol=lf +*.yaml text eol=lf +*.yml text eol=lf +*.md text eol=lf +*.txt text eol=lf +*.json text eol=lf +*.toml text eol=lf +*.sh text eol=lf +*.bat text eol=crlf +*.cmd text eol=crlf + +# Declare files that will always have CRLF line endings on checkout +*.bat text eol=crlf +*.cmd text eol=crlf + +# Denote all files that are truly binary and should not be modified +*.png binary +*.jpg binary +*.jpeg binary +*.gif binary +*.ico binary +*.mov binary +*.mp4 binary +*.mp3 binary +*.flv binary +*.fla binary +*.swf binary +*.gz binary +*.zip binary +*.7z binary +*.ttf binary +*.eot binary +*.woff binary +*.woff2 binary +*.pyc binary +*.pdf binary +*.ai binary +*.psd binary +*.xcf binary diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..93e843ae5 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,29 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + # list of supported hooks: https://pre-commit.com/hooks.html + - id: check-yaml + args: ["--unsafe"] + - id: debug-statements + - id: end-of-file-fixer + - id: mixed-line-ending + args: ["--fix=lf"] + - id: trailing-whitespace + args: ["--markdown-linebreak-ext=md"] + + - repo: https://github.com/PyCQA/isort + rev: 5.11.5 + hooks: + - id: isort + + - repo: https://github.com/psf/black + rev: 23.1.0 + hooks: + - id: black + + - repo: https://github.com/PyCQA/flake8 + rev: 4.0.1 + hooks: + - id: flake8 + args: ["--max-line-length=88", "--extend-ignore=E203,W503"] diff --git a/docs/api/API_DESIGN_TEMPLATE.md b/docs/api/API_DESIGN_TEMPLATE.md index eefaf9334..f9b0478dd 100644 --- a/docs/api/API_DESIGN_TEMPLATE.md +++ b/docs/api/API_DESIGN_TEMPLATE.md @@ -73,20 +73,20 @@ from enum import Enum class DiTCacheManager: """Manages DiT cache for diffusion models.""" - + def __init__(self, config: DiTCacheConfig): self.config = config self.cache = {} self.cache_stats = {} - + async def get_cache(self, cache_key: str) -> Optional[Any]: """Retrieve cached data.""" pass - + async def set_cache(self, cache_key: str, data: Any) -> None: """Store data in cache.""" pass - + async def invalidate_cache(self, cache_key: str) -> None: """Invalidate cached data.""" pass @@ -103,7 +103,7 @@ class OmniScheduler(): self.scheduler = self._create_scheduler() self.cache_manager = DiTCacheManager(config.dit_cache_config) self._running = False - + def _create_scheduler(self) -> BaseScheduler: """Factory method to create appropriate scheduler.""" if self.config.scheduler_type == SchedulerType.FIFO: @@ -118,13 +118,13 @@ class OmniScheduler(): async def schedule(self, request: OmniRequest) -> bool: """ Schedule a request for processing. - + Args: request: The request to schedule - + Returns: bool: True if successfully scheduled, False otherwise - + Raises: SchedulerError: If scheduling fails QueueFullError: If queue is at capacity @@ -137,7 +137,7 @@ class OmniScheduler(): def update_config(self, new_config: SchedulerConfig) -> None: """Update scheduler configuration.""" pass - + def get_config(self) -> SchedulerConfig: """Get current configuration.""" pass @@ -149,12 +149,12 @@ class OmniScheduler(): """Start the scheduler.""" self._running = True # Start background tasks - + async def stop(self) -> None: """Stop the scheduler gracefully.""" self._running = False # Cleanup resources - + async def shutdown(self) -> None: """Force shutdown the scheduler.""" # Immediate cleanup @@ -170,7 +170,7 @@ class OmniScheduler(): "processed_requests": self.scheduler.processed_count(), "cache_hit_rate": self.cache_manager.get_hit_rate() } - + def get_metrics(self) -> Dict[str, float]: """Get performance metrics.""" pass @@ -182,21 +182,21 @@ class OmniScheduler(): @dataclass class OmniConfig: """Configuration for the core module.""" - + # Scheduler configuration scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) - + # Cache configuration dit_cache: DiTCacheConfig = field(default_factory=DiTCacheConfig) - + # Resource limits max_memory_gb: float = 16.0 max_gpu_utilization: float = 0.8 - + # Timeouts request_timeout: int = 300 worker_timeout: int = 60 - + def validate(self) -> None: """Validate configuration parameters.""" if self.max_memory_gb <= 0: diff --git a/docs/api/README.md b/docs/api/README.md index 681bbba81..58490a80e 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -157,10 +157,10 @@ Each API template follows this standardized structure: # Clone the repository git clone https://github.com/hsliuustc0106/vllm-omni.git cd vllm-omni - + # Install dependencies pip install -r requirements-dev.txt - + # Set up pre-commit hooks pre-commit install ``` diff --git a/docs/architecture/high level arch design.md b/docs/architecture/high level arch design.md index 1545ca877..82a9f1a14 100644 --- a/docs/architecture/high level arch design.md +++ b/docs/architecture/high level arch design.md @@ -52,7 +52,7 @@ According to analysis for current popular open-source models, most of them have * **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 @@ -112,8 +112,8 @@ vLLM-omni will import core modules from existing vLLM framework and introduce se # Contact -For any questions or feedback regarding this design document, please contact -Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk), +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]: diff --git a/docs/architecture/implementation_architecture.md b/docs/architecture/implementation_architecture.md index a83d96d1e..2551885a8 100644 --- a/docs/architecture/implementation_architecture.md +++ b/docs/architecture/implementation_architecture.md @@ -136,7 +136,7 @@ class OmniLLM(LLM): self.engine_list: List[LLMEngine] = [] self.output_processor = MultimodalOutputProcessor() self._initialize_stage_engines() - + def _initialize_stage_engines(self) -> None: """Initialize LLMEngine instances for each stage""" for stage_config in self.stage_configs: @@ -145,24 +145,24 @@ class OmniLLM(LLM): elif stage_config.engine_type == "DiT": engine = self._create_dit_engine(stage_config) self.engine_list.append(engine) - + def generate(self, stage_args_list: List[Dict], **kwargs) -> List[RequestOutput]: """Main generation interface - orchestrates multi-stage processing""" current_output = None - + for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)): stage_engine = self.engine_list[i] - + # Prepare input for this stage processed_input = self._process_stage_inputs(stage_config, stage_args, current_output) - + # Execute stage stage_output = self._execute_stage(stage_engine, processed_input) - + # Update for next stage current_output = stage_output stage_config.stage_output = stage_output - + # Process final output return self.output_processor.process_output(current_output) ``` @@ -178,7 +178,7 @@ class AsyncOmniLLM(AsyncLLM): self.async_engine_list: List[AsyncLLM] = [] self.output_processor = MultimodalOutputProcessor() self._initialize_async_stage_engines() - + async def generate_async(self, stage_args_list: List[Dict], **kwargs) -> List[RequestOutput]: """Async generation interface""" # Similar to OmniLLM but with async/await patterns @@ -241,7 +241,7 @@ class MultimodalOutputProcessor(OutputProcessor): "latents": self._process_latents_output, "text": self._process_text_output, } - + def process_outputs(self, engine_core_outputs: List[EngineCoreOutput], ...): """Process multimodal outputs based on type""" for engine_core_output in engine_core_outputs: diff --git a/examples/basic/api_client.py b/examples/basic/api_client.py index 61e45a81b..19c1662aa 100644 --- a/examples/basic/api_client.py +++ b/examples/basic/api_client.py @@ -6,9 +6,9 @@ This example shows how to interact with the vLLM-omni API server using HTTP requests. """ -import requests import json -import time + +import requests def test_health(host="localhost", port=8000): @@ -30,7 +30,7 @@ def test_health(host="localhost", port=8000): def generate_text(prompt, host="localhost", port=8000, **kwargs): """Generate text using the API server.""" url = f"http://{host}:{port}/generate" - + # Default parameters params = { "prompts": [prompt], @@ -40,11 +40,11 @@ def generate_text(prompt, host="localhost", port=8000, **kwargs): "frequency_penalty": kwargs.get("frequency_penalty", 0.0), "presence_penalty": kwargs.get("presence_penalty", 0.0), } - + # Add optional parameters if "stop" in kwargs: params["stop"] = kwargs["stop"] - + try: response = requests.post(url, json=params, timeout=30) response.raise_for_status() @@ -58,71 +58,64 @@ def main(): """Main example function.""" print("vLLM-omni API Client Example") print("=" * 40) - + # Test server health if not test_health(): print("Please start the server first:") print(" vllm serve Qwen/Qwen3-0.6B --omni --port 8000") return - + # Example 1: Simple generation print("\n1. Simple text generation:") result = generate_text("Hello, how are you?", max_tokens=50) if result: text = result["outputs"][0]["text"] - print(f" Input: Hello, how are you?") + print(" Input: Hello, how are you?") print(f" Output: {text}") - + # Example 2: Creative writing print("\n2. Creative writing:") result = generate_text( "Write a short story about a robot learning to paint", max_tokens=150, - temperature=0.9 + temperature=0.9, ) if result: text = result["outputs"][0]["text"] - print(f" Input: Write a short story about a robot learning to paint") + print(" Input: Write a short story about a robot learning to paint") print(f" Output: {text}") - + # Example 3: Question answering print("\n3. Question answering:") result = generate_text( - "What is the capital of France?", - max_tokens=30, - temperature=0.3 + "What is the capital of France?", max_tokens=30, temperature=0.3 ) if result: text = result["outputs"][0]["text"] - print(f" Input: What is the capital of France?") + print(" Input: What is the capital of France?") print(f" Output: {text}") - + # Example 4: Multiple prompts print("\n4. Multiple prompts:") - prompts = [ - "Tell me a joke", - "What's 2+2?", - "Describe the color blue" - ] - + prompts = ["Tell me a joke", "What's 2+2?", "Describe the color blue"] + for prompt in prompts: result = generate_text(prompt, max_tokens=30) if result: text = result["outputs"][0]["text"] print(f" Q: {prompt}") print(f" A: {text}") - + # Example 5: Show full response structure print("\n5. Full response structure:") result = generate_text("Hello world", max_tokens=20) if result: print(" Full API response:") print(json.dumps(result, indent=2)) - + print("\n" + "=" * 40) print("API client examples completed!") if __name__ == "__main__": main() - diff --git a/examples/basic/text_generation.py b/examples/basic/text_generation.py index a6a9b5179..504f1c383 100644 --- a/examples/basic/text_generation.py +++ b/examples/basic/text_generation.py @@ -7,40 +7,43 @@ with a single AR (Autoregressive) stage. """ import asyncio -from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM + from vllm_omni.config import create_ar_stage_config +from vllm_omni.entrypoints.omni_llm import AsyncOmniLLM, OmniLLM def sync_example(): """Synchronous usage example.""" print("=== Synchronous vLLM-omni Example ===") - + # Create a simple AR stage configuration stage_config = create_ar_stage_config( stage_id=0, model_path="Qwen/Qwen3-0.6B", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) - + # Initialize OmniLLM omni_llm = OmniLLM([stage_config]) - + # Prepare stage arguments - stage_args = [{ - "prompt": "Hello, how are you today?", - "max_tokens": 50, - "temperature": 0.7 - }] - + stage_args = [ + { + "prompt": "Hello, how are you today?", + "max_tokens": 50, + "temperature": 0.7, + } + ] + # Generate text print(f"Input: {stage_args[0]['prompt']}") outputs = omni_llm.generate(stage_args) - + # Display results for i, output in enumerate(outputs): print(f"Output {i}:") - if hasattr(output, 'outputs') and output.outputs: + if hasattr(output, "outputs") and output.outputs: for completion in output.outputs: print(f" Text: {completion.text}") print(f" Finished: {completion.finish_reason != 'length'}") @@ -50,33 +53,35 @@ def sync_example(): async def async_example(): """Asynchronous usage example.""" print("\n=== Asynchronous vLLM-omni Example ===") - + # Create a simple AR stage configuration stage_config = create_ar_stage_config( stage_id=0, model_path="Qwen/Qwen3-0.6B", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) - + # Initialize AsyncOmniLLM omni_llm = AsyncOmniLLM([stage_config]) - + # Prepare stage arguments - stage_args = [{ - "prompt": "What is artificial intelligence?", - "max_tokens": 100, - "temperature": 0.8 - }] - + stage_args = [ + { + "prompt": "What is artificial intelligence?", + "max_tokens": 100, + "temperature": 0.8, + } + ] + # Generate text asynchronously print(f"Input: {stage_args[0]['prompt']}") outputs = await omni_llm.generate_async(stage_args) - + # Display results for i, output in enumerate(outputs): print(f"Output {i}:") - if hasattr(output, 'outputs') and output.outputs: + if hasattr(output, "outputs") and output.outputs: for completion in output.outputs: print(f" Text: {completion.text}") print(f" Finished: {completion.finish_reason != 'length'}") @@ -86,36 +91,32 @@ async def async_example(): def multi_prompt_example(): """Example with multiple prompts.""" print("\n=== Multi-Prompt Example ===") - + # Create stage configuration stage_config = create_ar_stage_config( stage_id=0, model_path="Qwen/Qwen3-0.6B", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) - + # Initialize OmniLLM omni_llm = OmniLLM([stage_config]) - + # Multiple prompts prompts = [ "Tell me a joke", "What's the weather like?", - "Explain quantum computing" + "Explain quantum computing", ] - + for prompt in prompts: print(f"\nInput: {prompt}") - stage_args = [{ - "prompt": prompt, - "max_tokens": 30, - "temperature": 0.9 - }] - + stage_args = [{"prompt": prompt, "max_tokens": 30, "temperature": 0.9}] + outputs = omni_llm.generate(stage_args) - - if outputs and hasattr(outputs[0], 'outputs') and outputs[0].outputs: + + if outputs and hasattr(outputs[0], "outputs") and outputs[0].outputs: completion = outputs[0].outputs[0] print(f"Output: {completion.text}") @@ -123,15 +124,15 @@ def multi_prompt_example(): if __name__ == "__main__": print("vLLM-omni Simple Usage Examples") print("=" * 40) - + # Run synchronous example sync_example() - + # Run asynchronous example asyncio.run(async_example()) - + # Run multi-prompt example multi_prompt_example() - + print("\n" + "=" * 40) print("Examples completed!") diff --git a/examples/omni/ar_dit_diffusers.py b/examples/omni/ar_dit_diffusers.py index 5685b5dc3..b5cfb176a 100644 --- a/examples/omni/ar_dit_diffusers.py +++ b/examples/omni/ar_dit_diffusers.py @@ -10,8 +10,8 @@ from typing import Dict, List import yaml from vllm_omni import ( - OmniLLM, DiTConfig, + OmniLLM, create_ar_stage_config, create_dit_stage_config, ) diff --git a/pyproject.toml b/pyproject.toml index 4e3382ccf..2219a7110 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,6 +103,10 @@ multi_line_output = 3 line_length = 88 known_first_party = ["vllm_omni"] +[tool.flake8] +max-line-length = 88 +extend-ignore = ["E203", "W503"] + [tool.mypy] python_version = "3.8" warn_return_any = true diff --git a/scripts/README.md b/scripts/README.md index 6f4e90bb5..ec097a9ed 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -106,4 +106,3 @@ python scripts/download_models.py - `server.log`: Server startup and runtime logs - `api_client_test.log`: API client test results - `simple_usage_test.log`: Simple usage test results - diff --git a/scripts/test_serving.sh b/scripts/test_serving.sh index d3a8d6cda..a9f1803cc 100755 --- a/scripts/test_serving.sh +++ b/scripts/test_serving.sh @@ -99,15 +99,15 @@ print('✅ All imports successful') start_server() { log_info "Starting vLLM-omni server on port $PORT..." log_info "Command: vllm serve $MODEL_PATH --omni --port $PORT" - + # Start server in background vllm serve "$MODEL_PATH" --omni --port "$PORT" > server.log 2>&1 & SERVER_PID=$! - + log_info "Server started with PID: $SERVER_PID" log_info "Waiting $SERVER_STARTUP_WAIT seconds for server to initialize..." sleep $SERVER_STARTUP_WAIT - + # Check if server is still running if ! kill -0 $SERVER_PID 2>/dev/null; then log_error "Server failed to start" @@ -115,7 +115,7 @@ start_server() { cat server.log exit 1 fi - + log_success "Server appears to be running" } @@ -146,7 +146,7 @@ test_info() { local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/info 2>/dev/null || echo "000") local http_code="${response: -3}" local body="${response%???}" - + if [ "$http_code" = "200" ]; then log_success "Info endpoint working" echo "$body" | python -m json.tool 2>/dev/null || log_warning "Info response not valid JSON" @@ -159,13 +159,13 @@ test_info() { # Test text generation test_generation() { log_info "Testing text generation..." - + # Test 1: Simple generation local response=$(curl -s -X POST http://$HOST:$PORT/generate \ -H "Content-Type: application/json" \ -d '{"prompts": ["Test the server functionality"], "max_tokens": 20, "temperature": 0.7}' \ 2>/dev/null || echo "{}") - + if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data and len(data['outputs']) > 0 else 1)" 2>/dev/null; then log_success "Text generation test passed" echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); print('Generated text:', data['outputs'][0]['text'][:100] + '...' if len(data['outputs'][0]['text']) > 100 else data['outputs'][0]['text'])" 2>/dev/null @@ -271,7 +271,7 @@ test_simple_usage() { # Performance test test_performance() { log_info "Running performance test..." - + local start_time=$(date +%s) local response=$(curl -s -X POST http://$HOST:$PORT/generate \ -H "Content-Type: application/json" \ @@ -279,7 +279,7 @@ test_performance() { 2>/dev/null) local end_time=$(date +%s) local duration=$((end_time - start_time)) - + if echo "$response" | python -c "import sys, json; data=json.load(sys.stdin); exit(0 if 'outputs' in data else 1)" 2>/dev/null; then log_success "Performance test passed (${duration}s response time)" else @@ -295,16 +295,16 @@ run_tests() { log_info "Port: $PORT" log_info "Host: $HOST" echo "==========================================" - + # Run all tests check_model check_environment test_imports start_server - + # Wait a bit more for server to be fully ready sleep 5 - + test_health test_info test_generation @@ -348,13 +348,13 @@ main() { show_usage exit 0 fi - + # Change to script directory cd "$(dirname "$0")/.." - + # Create logs directory if it doesn't exist mkdir -p logs - + # Run tests run_tests } diff --git a/tests/__init__.py b/tests/__init__.py index 0fc4af113..0a7176f94 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,5 +1,6 @@ """ Test suite for vLLM-omni. -This package contains unit tests, integration tests, and benchmarks 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 index cdec37201..3a857b8f4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,10 +2,17 @@ Shared fixtures and configuration for vLLM-omni tests. """ +from unittest.mock import Mock + import pytest import torch -from unittest.mock import Mock -from vllm_omni.config import OmniStageConfig, DiTConfig, DiTCacheConfig, create_ar_stage_config, create_dit_stage_config + +from vllm_omni.config import ( + DiTCacheConfig, + DiTConfig, + create_ar_stage_config, + create_dit_stage_config, +) @pytest.fixture(scope="session") @@ -21,7 +28,7 @@ def sample_ar_stage_config(): stage_id=0, model_path="test-ar-model", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) @@ -32,15 +39,15 @@ def sample_dit_stage_config(): model_type="dit", scheduler_type="ddpm", num_inference_steps=10, - guidance_scale=7.5 + guidance_scale=7.5, ) - + return create_dit_stage_config( stage_id=1, model_path="test-dit-model", input_modalities=["text"], output_modalities=["image"], - dit_config=dit_config + dit_config=dit_config, ) @@ -64,19 +71,16 @@ def mock_vllm_config(): def mock_dit_cache_config(): """Mock DiT cache configuration.""" from vllm_omni.config import DiTCacheTensor - + cache_tensors = [ DiTCacheTensor( - name="test_tensor", - shape=[1, 512, 512], - dtype="float32", - persistent=True + name="test_tensor", shape=[1, 512, 512], dtype="float32", persistent=True ) ] - + return DiTCacheConfig( cache_tensors=cache_tensors, max_cache_size=1024 * 1024 * 1024, # 1GB cache_strategy="fifo", - enable_optimization=True - ) \ No newline at end of file + enable_optimization=True, + ) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 1d8229156..14c6c8cfa 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -3,11 +3,12 @@ Unit tests for configuration modules. """ import pytest + from vllm_omni.config import ( - OmniStageConfig, - DiTConfig, DiTCacheConfig, DiTCacheTensor, + DiTConfig, + OmniStageConfig, create_ar_stage_config, create_dit_stage_config, ) @@ -21,33 +22,29 @@ class TestOmniStageConfig: engine_type="AR", model_path="test-model", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) assert config.engine_type == "AR" assert config.stage_id == 0 assert config.model_path == "test-model" assert config.input_modalities == ["text"] assert config.output_modalities == ["text"] - + def test_dit_stage_config_creation(self): """Test DiT stage configuration creation.""" - dit_config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=50 - ) + dit_config = DiTConfig(num_inference_steps=50) config = OmniStageConfig( stage_id=1, engine_type="DiT", model_path="test-dit-model", input_modalities=["text"], output_modalities=["image"], - dit_config=dit_config + dit_config=dit_config, ) assert config.engine_type == "DiT" assert config.dit_config is not None - assert config.dit_config.model_type == "dit" - + assert config.dit_config.num_inference_steps == 50 + def test_invalid_engine_type(self): """Test validation of engine type.""" with pytest.raises(ValueError): @@ -56,9 +53,9 @@ class TestOmniStageConfig: engine_type="INVALID", model_path="test-model", input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) - + def test_empty_modalities(self): """Test validation of empty modalities.""" with pytest.raises(ValueError): @@ -67,18 +64,18 @@ class TestOmniStageConfig: engine_type="AR", model_path="test-model", input_modalities=[], - output_modalities=["text"] + output_modalities=["text"], ) - + with pytest.raises(ValueError): OmniStageConfig( stage_id=0, engine_type="AR", model_path="test-model", input_modalities=["text"], - output_modalities=[] + output_modalities=[], ) - + def test_dit_requires_config(self): """Test that DiT engine requires dit_config.""" with pytest.raises(ValueError): @@ -87,32 +84,26 @@ class TestOmniStageConfig: engine_type="DiT", model_path="test-dit-model", input_modalities=["text"], - output_modalities=["image"] + output_modalities=["image"], ) class TestDiTConfig: def test_dit_config_defaults(self): """Test DiT configuration with defaults.""" - config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=50 - ) + config = DiTConfig(num_inference_steps=50) assert config.use_diffusers is False assert config.diffusers_pipeline is None assert config.guidance_scale == 7.5 assert config.height == 512 assert config.width == 512 - + def test_diffusers_config(self): """Test DiT configuration with diffusers.""" config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", num_inference_steps=50, use_diffusers=True, - diffusers_pipeline="stable-diffusion" + diffusers_pipeline="stable-diffusion", ) assert config.use_diffusers is True assert config.diffusers_pipeline == "stable-diffusion" @@ -126,16 +117,16 @@ class TestDiTCacheConfig: name="test_tensor", shape=[1, 512, 512], dtype="float32", - persistent=True + persistent=True, ) ] - + config = DiTCacheConfig( cache_tensors=cache_tensors, max_cache_size=1024 * 1024 * 1024, - cache_strategy="fifo" + cache_strategy="fifo", ) - + assert len(config.cache_tensors) == 1 assert config.max_cache_size == 1024 * 1024 * 1024 assert config.cache_strategy == "fifo" @@ -145,35 +136,25 @@ class TestDiTCacheConfig: class TestHelperFunctions: def test_create_ar_stage_config(self): """Test create_ar_stage_config helper function.""" - config = create_ar_stage_config( - stage_id=0, - model_path="test-model" - ) - + config = create_ar_stage_config(stage_id=0, model_path="test-model") + assert config.stage_id == 0 assert config.engine_type == "AR" assert config.model_path == "test-model" assert config.input_modalities == ["text"] assert config.output_modalities == ["text"] - + def test_create_dit_stage_config(self): """Test create_dit_stage_config helper function.""" - dit_config = DiTConfig( - model_type="dit", - scheduler_type="ddpm", - num_inference_steps=50 - ) - + dit_config = DiTConfig(num_inference_steps=50) + config = create_dit_stage_config( - stage_id=1, - model_path="test-dit-model", - dit_config=dit_config + stage_id=1, model_path="test-dit-model", dit_config=dit_config ) - + assert config.stage_id == 1 assert config.engine_type == "DiT" assert config.model_path == "test-dit-model" assert config.input_modalities == ["text"] assert config.output_modalities == ["image"] assert config.dit_config is not None - diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index fb544b0fa..e0b6d0e4a 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -1,45 +1,46 @@ """ -vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures. +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. +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 +- 🟡 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 .entrypoints.omni_llm import OmniLLM, AsyncOmniLLM from .config import ( - OmniStageConfig, - DiTConfig, DiTCacheConfig, + DiTConfig, + OmniStageConfig, create_ar_stage_config, create_dit_stage_config, ) +# Main entry points +from .entrypoints.omni_llm import AsyncOmniLLM, OmniLLM + __all__ = [ # Version info "__version__", - "__author__", + "__author__", "__email__", - # Main components "OmniLLM", "AsyncOmniLLM", - # Configuration "OmniStageConfig", "DiTConfig", "DiTCacheConfig", "create_ar_stage_config", "create_dit_stage_config", - # 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 index a688faf1b..2b4a732f1 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -11,10 +11,10 @@ from vllm.config import config import vllm_omni.model_executor.models as me_models from .stage_config import ( - OmniStageConfig, - DiTConfig, DiTCacheConfig, DiTCacheTensor, + DiTConfig, + OmniStageConfig, create_ar_stage_config, create_dit_stage_config, ) @@ -42,9 +42,9 @@ class OmniModelConfig(ModelConfig): __all__ = [ "OmniModelConfig", "OmniStageConfig", - "DiTConfig", + "DiTConfig", "DiTCacheConfig", "DiTCacheTensor", "create_ar_stage_config", "create_dit_stage_config", -] \ No newline at end of file +] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 8cca98c63..e8156e7dd 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -3,19 +3,22 @@ Stage configuration for vLLM-omni multi-stage processing. """ from dataclasses import dataclass -from typing import List, Optional, Type, Literal, Any, Dict +from typing import Any, Dict, List, Literal, Optional, Type + from vllm.config import ( - VllmConfig, + CompilationConfig, DeviceConfig, LoadConfig, ModelConfig, - CompilationConfig, + VllmConfig, ) from vllm.executor.executor_base import ExecutorBase as Executor + @dataclass class DiTCacheTensor: """Configuration for DiT cache tensors.""" + name: str shape: List[int] dtype: str = "float32" @@ -25,6 +28,7 @@ class DiTCacheTensor: @dataclass class DiTCacheConfig: """Configuration for DiT caching system.""" + cache_tensors: List[DiTCacheTensor] max_cache_size: int = 1024 * 1024 * 1024 # 1GB cache_strategy: str = "fifo" # fifo, lru, lfu @@ -73,38 +77,47 @@ class OmniStageConfig: executor_class: Type[Executor] = None # Will be set based on engine_type default_stage_args: Optional[Dict[str, Any]] = None stage_output: Optional[Any] = None - + def __post_init__(self): """Validate configuration after initialization.""" if self.engine_type not in ["AR", "DiT"]: - raise ValueError(f"Invalid engine_type: {self.engine_type}. Must be 'AR' or 'DiT'") - + raise ValueError( + f"Invalid engine_type: {self.engine_type}. Must be 'AR' or 'DiT'" + ) + if self.engine_type == "DiT" and self.dit_config is None: raise ValueError("DiT engine requires dit_config") - + if not self.input_modalities: raise ValueError("input_modalities cannot be empty") - + if not self.output_modalities: raise ValueError("output_modalities cannot be empty") - + def get_executor_class(self) -> Type[Executor]: """Get the appropriate executor class for this stage.""" if self.executor_class is not None: return self.executor_class - + if self.engine_type == "AR": from vllm.executor.uniproc_executor import UniProcExecutor + return UniProcExecutor elif self.engine_type == "DiT": if self.dit_config and self.dit_config.use_diffusers: - from vllm_omni.executor.diffusers_executor import DiffusersPipelineExecutor + from vllm_omni.executor.diffusers_executor import ( + DiffusersPipelineExecutor, + ) + return DiffusersPipelineExecutor else: from vllm.executor.uniproc_executor import UniProcExecutor + return UniProcExecutor - - raise ValueError(f"No executor class available for engine_type: {self.engine_type}") + + raise ValueError( + f"No executor class available for engine_type: {self.engine_type}" + ) def create_ar_stage_config( @@ -121,7 +134,7 @@ def create_ar_stage_config( input_modalities = ["text"] if output_modalities is None: output_modalities = ["text"] - + # For now, we'll create minimal configs # vllm_config and executor_class will be handled by the LLM class return OmniStageConfig( @@ -150,7 +163,7 @@ def create_dit_stage_config( input_modalities = ["text"] if output_modalities is None: output_modalities = ["image"] - + if dit_config is None: dit_config = DiTConfig(num_inference_steps=50) diff --git a/vllm_omni/core/dit_cache_manager.py b/vllm_omni/core/dit_cache_manager.py index d2c800e76..4bdeec47b 100644 --- a/vllm_omni/core/dit_cache_manager.py +++ b/vllm_omni/core/dit_cache_manager.py @@ -1,20 +1,21 @@ """ DiT Cache Manager for vLLM-omni. -This module provides caching functionality for DiT (Diffusion Transformer) models -to optimize inference performance and memory usage. +This module provides caching functionality for DiT (Diffusion Transformer) +models to optimize inference performance and memory usage. """ import time +from typing import Any, Dict, List, Optional + import torch -from typing import Dict, List, Optional, Any -from dataclasses import dataclass + from ..config import DiTCacheConfig, DiTCacheTensor class DiTCacheManager: """Manages DiT-specific caching for optimized inference.""" - + def __init__(self, config: DiTCacheConfig): self.config = config self.cache_tensors: Dict[str, DiTCacheTensor] = {} @@ -24,30 +25,30 @@ class DiTCacheManager: self.cache_hits = 0 self.cache_misses = 0 self.cache_timeout = 3600 # 1 hour default timeout - + def allocate_cache(self, request_id: str, size: int) -> torch.Tensor: """Allocate cache for a specific request.""" # Check if we have enough space if self.current_cache_size + size > self.max_cache_size: self._evict_cache() - + # Create tensor tensor = torch.zeros(size, dtype=torch.float32) - + # Store in cache cache_tensor = DiTCacheTensor( name=request_id, tensor=tensor, timestamp=time.time(), access_count=0, - size_bytes=size * 4 # Assuming float32 + size_bytes=size * 4, # Assuming float32 ) - + self.cache_tensors[request_id] = cache_tensor self.current_cache_size += cache_tensor.size_bytes - + return tensor - + def get_cache(self, request_id: str) -> Optional[torch.Tensor]: """Get cached tensor for a request.""" if request_id in self.cache_tensors: @@ -58,58 +59,58 @@ class DiTCacheManager: else: self.cache_misses += 1 return None - + def store_cache(self, request_id: str, tensor: torch.Tensor) -> None: """Store a tensor in the cache.""" if request_id in self.cache_tensors: # Update existing cache old_tensor = self.cache_tensors[request_id] self.current_cache_size -= old_tensor.size_bytes - + new_size_bytes = tensor.numel() * 4 # Assuming float32 self.current_cache_size += new_size_bytes - + self.cache_tensors[request_id] = DiTCacheTensor( name=request_id, tensor=tensor, timestamp=time.time(), access_count=old_tensor.access_count, - size_bytes=new_size_bytes + size_bytes=new_size_bytes, ) else: # Create new cache entry new_size_bytes = tensor.numel() * 4 if self.current_cache_size + new_size_bytes > self.max_cache_size: self._evict_cache() - + self.cache_tensors[request_id] = DiTCacheTensor( name=request_id, tensor=tensor, timestamp=time.time(), access_count=0, - size_bytes=new_size_bytes + size_bytes=new_size_bytes, ) self.current_cache_size += new_size_bytes - + def release_cache(self, request_id: str) -> None: """Release cache for a request.""" if request_id in self.cache_tensors: cache_tensor = self.cache_tensors[request_id] self.current_cache_size -= cache_tensor.size_bytes del self.cache_tensors[request_id] - + def clear_expired_cache(self) -> None: """Clear expired cache entries.""" current_time = time.time() expired_keys = [] - + for request_id, cache_tensor in self.cache_tensors.items(): if current_time - cache_tensor.timestamp > self.cache_timeout: expired_keys.append(request_id) - + for request_id in expired_keys: self.release_cache(request_id) - + def _evict_cache(self) -> None: """Evict cache entries based on strategy.""" if self.config.cache_strategy == "fifo": @@ -120,42 +121,48 @@ class DiTCacheManager: self._evict_lfu() else: self._evict_fifo() # Default to FIFO - + def _evict_fifo(self) -> None: """Evict cache entries using FIFO strategy.""" if not self.cache_tensors: return - + # Find oldest entry - oldest_key = min(self.cache_tensors.keys(), - key=lambda k: self.cache_tensors[k].timestamp) + oldest_key = min( + self.cache_tensors.keys(), + key=lambda k: self.cache_tensors[k].timestamp, + ) self.release_cache(oldest_key) - + def _evict_lru(self) -> None: """Evict cache entries using LRU strategy.""" if not self.cache_tensors: return - + # Find least recently used entry - lru_key = min(self.cache_tensors.keys(), - key=lambda k: self.cache_tensors[k].access_count) + lru_key = min( + self.cache_tensors.keys(), + key=lambda k: self.cache_tensors[k].access_count, + ) self.release_cache(lru_key) - + def _evict_lfu(self) -> None: """Evict cache entries using LFU strategy.""" if not self.cache_tensors: return - + # Find least frequently used entry - lfu_key = min(self.cache_tensors.keys(), - key=lambda k: self.cache_tensors[k].access_count) + lfu_key = min( + self.cache_tensors.keys(), + key=lambda k: self.cache_tensors[k].access_count, + ) self.release_cache(lfu_key) - + def get_statistics(self) -> Dict[str, Any]: """Get cache statistics.""" total_requests = self.cache_hits + self.cache_misses hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0 - + return { "cache_size": self.current_cache_size, "max_cache_size": self.max_cache_size, @@ -163,16 +170,16 @@ class DiTCacheManager: "cache_hits": self.cache_hits, "cache_misses": self.cache_misses, "hit_rate": hit_rate, - "num_cached_tensors": len(self.cache_tensors) + "num_cached_tensors": len(self.cache_tensors), } - + def clear_all_cache(self) -> None: """Clear all cache entries.""" self.cache_tensors.clear() self.current_cache_size = 0 self.cache_hits = 0 self.cache_misses = 0 - + def set_cache_timeout(self, timeout: float) -> None: """Set cache timeout in seconds.""" - self.cache_timeout = timeout \ No newline at end of file + self.cache_timeout = timeout diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index c2ab49b6b..40e5b34a1 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -7,8 +7,9 @@ class OmniScheduler(VLLMScheduler): """ OmniScheduler: Scheduler for vLLM-omni multimodal processing. - This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive - processing with additional fields and methods specific to vLLM-omni. + This scheduler extends vLLM's scheduler to support multimodal and + non-autoregressive processing with additional fields and methods + specific to vLLM-omni. """ # Ensure scheduled_new_reqs carry omni-specific payloads (e.g., additional_information) diff --git a/vllm_omni/engine/diffusion_engine.py b/vllm_omni/engine/diffusion_engine.py index 9009c98f5..3989efec7 100644 --- a/vllm_omni/engine/diffusion_engine.py +++ b/vllm_omni/engine/diffusion_engine.py @@ -5,10 +5,7 @@ from __future__ import annotations from typing import Any, Optional from ..config import DiTConfig -from ..worker.gpu_diffusion_worker import ( - DiffusionGPUWorker, - DiffusionRunnerOutput, -) +from ..worker.gpu_diffusion_worker import DiffusionGPUWorker, DiffusionRunnerOutput class DiffusionEngine: @@ -27,14 +24,12 @@ class DiffusionEngine: config_model_path = None if model_cfg is not None: if isinstance(model_cfg, dict): - config_model_path = ( - model_cfg.get("model") - or model_cfg.get("model_path") + config_model_path = model_cfg.get("model") or model_cfg.get( + "model_path" ) else: - config_model_path = ( - getattr(model_cfg, "model", None) - or getattr(model_cfg, "model_path", None) + config_model_path = getattr(model_cfg, "model", None) or getattr( + model_cfg, "model_path", None ) self.model_path = model_path or config_model_path @@ -116,9 +111,7 @@ class DiffusersPipelineEngine(DiffusionEngine): else self.num_inference_steps ) guidance = ( - float(guidance_scale) - if guidance_scale is not None - else self.guidance_scale + float(guidance_scale) if guidance_scale is not None else self.guidance_scale ) return self.worker.generate( diff --git a/vllm_omni/entrypoints/api_server.py b/vllm_omni/entrypoints/api_server.py index af14fbf08..8bfb48067 100644 --- a/vllm_omni/entrypoints/api_server.py +++ b/vllm_omni/entrypoints/api_server.py @@ -3,17 +3,18 @@ API server for vLLM-omni. """ import asyncio -from typing import Dict, Any, List -from fastapi import FastAPI, HTTPException -from fastapi.responses import JSONResponse -from pydantic import BaseModel +from typing import Any, Dict, List + import uvicorn +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel from .omni_llm import AsyncOmniLLM class GenerateRequest(BaseModel): """Request model for generation.""" + prompts: List[str] max_tokens: int = 100 temperature: float = 0.7 @@ -22,6 +23,7 @@ class GenerateRequest(BaseModel): class GenerateResponse(BaseModel): """Response model for generation.""" + outputs: List[Dict[str, Any]] stage_outputs: List[Dict[str, Any]] = None @@ -29,7 +31,7 @@ class GenerateResponse(BaseModel): app = FastAPI( title="vLLM-omni API", description="Multi-modality models inference and serving", - version="0.1.0" + version="0.1.0", ) # Global omni_llm instance @@ -50,48 +52,51 @@ async def generate(request: GenerateRequest): try: if omni_llm is None: raise HTTPException(status_code=500, detail="OmniLLM not initialized") - + # Prepare stage arguments if request.stage_args is None: # Create default stage arguments - one per stage config # For now, we'll process all prompts in the first stage - stage_args = [{ - "prompt": " ".join(request.prompts) if request.prompts else "", - "max_tokens": request.max_tokens, - "temperature": request.temperature - }] + stage_args = [ + { + "prompt": " ".join(request.prompts) if request.prompts else "", + "max_tokens": request.max_tokens, + "temperature": request.temperature, + } + ] else: stage_args = request.stage_args - + # Generate using omni_llm try: outputs = await omni_llm.generate_async(stage_args) - except Exception as e: + except Exception: # Fallback to synchronous generation outputs = omni_llm.generate(stage_args) - + # Convert outputs to response format response_outputs = [] for output in outputs: - if hasattr(output, 'outputs') and output.outputs: + if hasattr(output, "outputs") and output.outputs: for out in output.outputs: - response_outputs.append({ - "text": getattr(out, 'text', ''), - "finished": getattr(out, 'finish_reason', 'length') != 'length', - "tokens": getattr(out, 'token_ids', []) - }) + response_outputs.append( + { + "text": getattr(out, "text", ""), + "finished": getattr(out, "finish_reason", "length") + != "length", + "tokens": getattr(out, "token_ids", []), + } + ) else: - response_outputs.append({ - "text": "", - "finished": True, - "tokens": [] - }) - + response_outputs.append({"text": "", "finished": True, "tokens": []}) + return GenerateResponse( outputs=response_outputs, - stage_outputs=[{"stage": i, "output": "processed"} for i in range(len(stage_args))] + stage_outputs=[ + {"stage": i, "output": "processed"} for i in range(len(stage_args)) + ], ) - + except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @@ -107,36 +112,35 @@ async def get_info(): """Get information about the service.""" if omni_llm is None: return {"error": "OmniLLM not initialized"} - + return { "service": "vllm-omni", "version": "0.1.0", - "num_stages": omni_llm.get_num_stages() if hasattr(omni_llm, 'get_num_stages') else 0, + "num_stages": omni_llm.get_num_stages() + if hasattr(omni_llm, "get_num_stages") + else 0, "stage_configs": [ { "stage_id": config.stage_id, "engine_type": config.engine_type, "model_path": config.model_path, "input_modalities": config.input_modalities, - "output_modalities": config.output_modalities + "output_modalities": config.output_modalities, } for config in omni_llm.stage_configs - ] + ], } -async def run_server(omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000): +async def run_server( + omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000 +): """Run the API server.""" global omni_llm omni_llm = omni_llm_instance - - config = uvicorn.Config( - app=app, - host=host, - port=port, - log_level="info" - ) - + + config = uvicorn.Config(app=app, host=host, port=port, log_level="info") + server = uvicorn.Server(config) await server.serve() diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 830e181a9..18d5ce9b6 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -3,8 +3,7 @@ CLI entry point for vLLM-omni that intercepts vLLM commands. """ import sys -import argparse -from typing import List, Optional + from vllm_omni.entrypoints.cli.serve import OmniServeCommand @@ -19,6 +18,7 @@ def main(): else: # Forward to original vLLM CLI from vllm.entrypoints.cli.main import main as vllm_main + vllm_main() diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 1fdc0f79a..0ee6caf35 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -4,125 +4,111 @@ Omni serve command for vLLM-omni. import argparse import asyncio -from typing import List, Optional +from typing import List + +from vllm_omni.config import DiTConfig, create_ar_stage_config, create_dit_stage_config from ..omni_llm import AsyncOmniLLM -from vllm_omni.config import ( - DiTConfig, - create_ar_stage_config, - create_dit_stage_config, -) class OmniServeCommand: """Command handler for vLLM-omni serve command.""" - + def __init__(self): self.parser = self._create_parser() - + def _create_parser(self) -> argparse.ArgumentParser: """Create argument parser for omni serve command.""" parser = argparse.ArgumentParser( description="vLLM-omni: Multi-modality models inference and serving" ) - + # Model arguments - make it optional with default parser.add_argument( "model", nargs="?", default="Qwen/Qwen3-0.6B", - help="Path to the model or model name (default: Qwen/Qwen3-0.6B)" + help="Path to the model or model name (default: Qwen/Qwen3-0.6B)", ) - + # Server arguments parser.add_argument( "--port", type=int, default=8000, - help="Port to run the server on" + help="Port to run the server on", ) - + parser.add_argument( "--host", type=str, default="0.0.0.0", - help="Host to run the server on" + help="Host to run the server on", ) - + # Stage configuration arguments - parser.add_argument( - "--ar-stage", - type=str, - help="AR stage model path" - ) - - parser.add_argument( - "--dit-stage", - type=str, - help="DiT stage model path" - ) - + parser.add_argument("--ar-stage", type=str, help="AR stage model path") + + parser.add_argument("--dit-stage", type=str, help="DiT stage model path") + parser.add_argument( "--dit-steps", type=int, default=50, - help="Number of DiT inference steps" + help="Number of DiT inference steps", ) - + parser.add_argument( "--dit-guidance-scale", type=float, default=7.5, - help="DiT guidance scale" + help="DiT guidance scale", ) - + parser.add_argument( "--use-diffusers", action="store_true", - help="Use diffusers pipeline for DiT stage" + help="Use diffusers pipeline for DiT stage", ) - + # Other arguments parser.add_argument( - "--log-stats", - action="store_true", - help="Enable logging statistics" + "--log-stats", action="store_true", help="Enable logging statistics" ) - + return parser - + def run(self, args: List[str]) -> None: """Run the omni serve command.""" parsed_args = self.parser.parse_args(args) - + # Create stage configurations stage_configs = self._create_stage_configs(parsed_args) - + # Create AsyncOmniLLM instance omni_llm = AsyncOmniLLM( - stage_configs=stage_configs, - log_stats=parsed_args.log_stats + stage_configs=stage_configs, log_stats=parsed_args.log_stats ) - + # Start the server asyncio.run(self._start_server(omni_llm, parsed_args)) - + def _create_stage_configs(self, args) -> List: """Create stage configurations based on arguments.""" stage_configs = [] stage_id = 0 - + # Add AR stage - use main model if ar_stage not specified ar_model = args.ar_stage if args.ar_stage else args.model ar_config = create_ar_stage_config( stage_id=stage_id, model_path=ar_model, input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) stage_configs.append(ar_config) stage_id += 1 - + # Add DiT stage if specified if args.dit_stage: dit_model = args.dit_stage @@ -131,7 +117,7 @@ class OmniServeCommand: dit_model = "stabilityai/stable-diffusion-2-1" else: dit_model = None - + if dit_model: dit_config = DiTConfig( num_inference_steps=args.dit_steps, @@ -139,17 +125,17 @@ class OmniServeCommand: use_diffusers=args.use_diffusers, diffusers_pipeline="auto" if args.use_diffusers else None, ) - + dit_stage_config = create_dit_stage_config( stage_id=stage_id, model_path=dit_model, input_modalities=["text"], output_modalities=["image"], - dit_config=dit_config + dit_config=dit_config, ) stage_configs.append(dit_stage_config) stage_id += 1 - + # If no specific stages are specified, use the main model if not stage_configs: # Try to detect if it's a multimodal model @@ -159,23 +145,23 @@ class OmniServeCommand: stage_id=0, model_path=args.model, input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) stage_configs.append(ar_config) - + dit_config = DiTConfig( num_inference_steps=args.dit_steps, guidance_scale=args.dit_guidance_scale, use_diffusers=args.use_diffusers, diffusers_pipeline="auto" if args.use_diffusers else None, ) - + dit_stage_config = create_dit_stage_config( stage_id=1, model_path=args.model, input_modalities=["text"], output_modalities=["image"], - dit_config=dit_config + dit_config=dit_config, ) stage_configs.append(dit_stage_config) else: @@ -184,23 +170,19 @@ class OmniServeCommand: stage_id=0, model_path=args.model, input_modalities=["text"], - output_modalities=["text"] + output_modalities=["text"], ) stage_configs.append(ar_config) - + return stage_configs - + async def _start_server(self, omni_llm: AsyncOmniLLM, args) -> None: """Start the API server.""" try: # Import here to avoid circular imports from vllm_omni.entrypoints.api_server import run_server - - await run_server( - omni_llm_instance=omni_llm, - host=args.host, - port=args.port - ) + + await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port) except KeyboardInterrupt: print("\nShutting down server...") except Exception as e: