diff --git a/.gitignore b/.gitignore index d42c87793..5b3538926 100644 --- a/.gitignore +++ b/.gitignore @@ -152,6 +152,12 @@ dmypy.json # Cython debug symbols cython_debug/ +# cursor +.cursor/ + +# docker +docker/ + # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be added to the global gitignore or merged into this project gitignore. For a PyCharm diff --git a/README.md b/README.md index 13da76012..66c8725ee 100644 --- a/README.md +++ b/README.md @@ -47,17 +47,100 @@ vLLM-omni is built on a modular architecture that extends vLLM's core functional ## ๐Ÿ› ๏ธ Installation -### Prerequisites +### Quick Start -- Python 3.8+ -- PyTorch 2.0+ -- CUDA 11.8+ (for GPU acceleration) - -### Install from Source +#### Option 1: Docker (Recommended for macOS) ```bash -git clone https://github.com/your-org/vllm-omni.git -pip install -r requirements.txt +# Clone the repository +git clone https://github.com/hsliuustc0106/vllm-omni.git cd vllm-omni + +# Run the automated Docker setup +./scripts/docker-setup-macos.sh +``` + +#### Option 2: Local Installation + +```bash +# Clone the repository +git clone https://github.com/hsliuustc0106/vllm-omni.git +cd vllm-omni + +# Run the installation script +./install.sh +``` + +### Prerequisites + +- Python 3.11+ (recommended) +- Conda or Miniconda +- Git +- CUDA 11.8+ (for GPU acceleration) or CPU-only installation + +### Installation Methods + +#### Method 1: Automated Installation (Recommended) +```bash +# Using shell script +./install.sh + +# Or using Python script +python install.py +``` + +#### Method 2: Manual Installation +```bash +# Create conda environment +conda create -n vllm_omni python=3.11 -y +conda activate vllm_omni + +# Install PyTorch (CPU or GPU) +pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cpu # CPU +# pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cu121 # GPU + +# Install dependencies +pip install -r requirements.txt +pip install "vllm>=0.10.2" + +# Install vLLM-omni pip install -e . ``` + +### Verify Installation + +```bash +# Test the installation +python test_installation.py + +# Test basic functionality +python -c "import vllm_omni; print('Ready!')" + +# Test CLI +vllm --help +``` + +For detailed installation instructions, see [INSTALL.md](INSTALL.md). + +## ๐Ÿ“ฅ Model Download + +Models are automatically downloaded when first used, or you can pre-download them: + +```bash +# Check downloaded models +python scripts/download_models.py --check-cache + +# Download all default models +python scripts/download_models.py --all + +# Download specific models +python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B +python scripts/download_models.py --dit-models stabilityai/stable-diffusion-2-1 +``` + +**Model Storage Location:** +- Default: `~/.cache/huggingface/hub/` +- AR models: 100MB - 1GB each +- DiT models: 2GB - 7GB each + +For detailed model management, see [MODEL_DOWNLOAD_GUIDE.md](docs/MODEL_DOWNLOAD_GUIDE.md). diff --git a/docs/implementation_summary.md b/docs/implementation_summary.md index 204a798c8..cdd3b4b45 100644 --- a/docs/implementation_summary.md +++ b/docs/implementation_summary.md @@ -218,7 +218,7 @@ from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config # Create stage configurations ar_config = create_ar_stage_config( stage_id=0, - model_path="microsoft/DialoGPT-small", + model_path="Qwen/Qwen3-0.6B", input_modalities=["text"], output_modalities=["text"] ) diff --git a/docs/testing/test_design.md b/docs/testing/test_design.md index 6206c1a35..7710dc0e0 100644 --- a/docs/testing/test_design.md +++ b/docs/testing/test_design.md @@ -373,7 +373,7 @@ class TestVLLMIntegration: OmniStageConfig( stage_id=0, engine_type="AR", - model_path="microsoft/DialoGPT-small", # Small test model + model_path="Qwen/Qwen3-0.6B", # Small test model input_modalities=["text"], output_modalities=["text"] ) @@ -443,7 +443,7 @@ class TestFullPipeline: OmniStageConfig( stage_id=0, engine_type="AR", - model_path="microsoft/DialoGPT-small", + model_path="Qwen/Qwen3-0.6B", input_modalities=["text"], output_modalities=["text"] ), diff --git a/examples/README.md b/examples/README.md index e79470c89..f09a1749c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,62 +1,49 @@ # vLLM-omni Examples -This directory contains examples demonstrating how to use vLLM-omni for various tasks. +This directory contains examples demonstrating how to use vLLM-omni for multi-modality model inference and serving. -## Basic Examples +## Directory Structure -- [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 +- `basic/` - Simple usage examples and API client +- `advanced/` - Advanced features and configurations +- `multimodal/` - Multi-modality specific examples -## Advanced Examples +## Quick Start -- [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: +1. **Basic Usage:** ```bash - pip install vllm-omni + cd examples/basic + python simple_usage.py ``` -2. Run a basic example: +2. **API Server:** ```bash - python examples/basic/text_generation.py + # Start server + vllm serve Qwen/Qwen3-0.6B --omni --port 8000 + + # Test with client + python api_client.py ``` -3. Explore the examples in each subdirectory for more advanced usage. +## Examples Overview + +### Basic Examples (`basic/`) +- `simple_usage.py` - Direct library usage (sync/async) +- `api_client.py` - HTTP API client examples +- `README.md` - Basic examples documentation + +### Advanced Examples (`advanced/`) +- Multi-stage processing +- Custom configurations +- Performance optimization + +### Multi-modality Examples (`multimodal/`) +- Text-to-image generation +- Image-to-text processing +- Multi-modal pipelines ## 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 -``` +- Python 3.8+ +- vLLM-omni installed +- For API examples: HTTP server running \ No newline at end of file diff --git a/examples/basic/README.md b/examples/basic/README.md new file mode 100644 index 000000000..6a0257fbb --- /dev/null +++ b/examples/basic/README.md @@ -0,0 +1,76 @@ +# Basic vLLM-omni Examples + +This directory contains simple examples demonstrating how to use vLLM-omni for text generation. + +## Examples + +### 1. `simple_usage.py` - Direct Library Usage + +Demonstrates how to use vLLM-omni directly in your Python code: + +```bash +# Run the example +python simple_usage.py +``` + +**Features:** +- Synchronous text generation +- Asynchronous text generation +- Multi-prompt processing +- Shows both text output and token information + +### 2. `api_client.py` - API Server Usage + +Demonstrates how to interact with the vLLM-omni API server: + +```bash +# First, start the server +vllm serve Qwen/Qwen3-0.6B --omni --port 8000 + +# Then run the client (in another terminal) +python api_client.py +``` + +**Features:** +- Health check +- Text generation via HTTP API +- Multiple prompt types +- Full response structure inspection + +## Quick Start + +1. **Install vLLM-omni:** + ```bash + pip install -e . + ``` + +2. **Run direct usage example:** + ```bash + cd examples/basic + python simple_usage.py + ``` + +3. **Run API server example:** + ```bash + # Terminal 1: Start server + vllm serve Qwen/Qwen3-0.6B --omni --port 8000 + + # Terminal 2: Run client + cd examples/basic + python api_client.py + ``` + +## Expected Output + +The examples will show: +- โœ… Generated text responses +- โœ… Token counts and completion status +- โœ… Stage processing information +- โœ… Error handling + +## Requirements + +- Python 3.8+ +- vLLM-omni installed +- For API examples: HTTP server running + diff --git a/examples/basic/api_client.py b/examples/basic/api_client.py new file mode 100644 index 000000000..61e45a81b --- /dev/null +++ b/examples/basic/api_client.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +""" +Simple API client example for vLLM-omni server. + +This example shows how to interact with the vLLM-omni API server +using HTTP requests. +""" + +import requests +import json +import time + + +def test_health(host="localhost", port=8000): + """Test if the server is healthy.""" + url = f"http://{host}:{port}/health" + try: + response = requests.get(url, timeout=5) + if response.status_code == 200: + print("โœ… Server is healthy") + return True + else: + print(f"โŒ Server health check failed: {response.status_code}") + return False + except requests.exceptions.RequestException as e: + print(f"โŒ Cannot connect to server: {e}") + return False + + +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], + "max_tokens": kwargs.get("max_tokens", 100), + "temperature": kwargs.get("temperature", 0.7), + "top_p": kwargs.get("top_p", 1.0), + "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() + return response.json() + except requests.exceptions.RequestException as e: + print(f"โŒ API request failed: {e}") + return None + + +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(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 + ) + if result: + text = result["outputs"][0]["text"] + print(f" 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 + ) + if result: + text = result["outputs"][0]["text"] + print(f" 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" + ] + + 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/docker_example.py b/examples/basic/docker_example.py new file mode 100644 index 000000000..2af59d863 --- /dev/null +++ b/examples/basic/docker_example.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +""" +Docker example for vLLM-omni. + +This example shows how to test the vLLM-omni API server running in Docker. +""" + +import requests +import json +import time +import subprocess +import sys + + +def check_docker_running(): + """Check if Docker is running.""" + try: + result = subprocess.run(['docker', 'info'], + capture_output=True, text=True, timeout=5) + return result.returncode == 0 + except (subprocess.TimeoutExpired, FileNotFoundError): + return False + + +def start_docker_server(): + """Start the vLLM-omni server in Docker.""" + print("๐Ÿณ Starting vLLM-omni server in Docker...") + + # Check if container is already running + result = subprocess.run(['docker', 'ps', '--filter', 'name=vllm-omni-cpu', '--format', '{{.Names}}'], + capture_output=True, text=True) + + if 'vllm-omni-cpu' in result.stdout: + print("โœ… Server already running") + return True + + # Start the server + cmd = [ + 'docker', 'run', '-d', '--name', 'vllm-omni-cpu', + '-p', '8000:8000', + 'vllm-omni-cpu:latest', + 'vllm', 'serve', 'Qwen/Qwen3-0.6B', '--omni', '--port', '8000' + ] + + try: + subprocess.run(cmd, check=True) + print("โœ… Server started successfully") + + # Wait for server to be ready + print("โณ Waiting for server to be ready...") + for i in range(30): # Wait up to 30 seconds + if test_server_health(): + print("โœ… Server is ready!") + return True + time.sleep(1) + + print("โŒ Server failed to start within 30 seconds") + return False + + except subprocess.CalledProcessError as e: + print(f"โŒ Failed to start server: {e}") + return False + + +def test_server_health(): + """Test if the server is healthy.""" + try: + response = requests.get("http://localhost:8000/health", timeout=2) + return response.status_code == 200 + except requests.exceptions.RequestException: + return False + + +def test_generation(): + """Test text generation via API.""" + print("\n๐Ÿงช Testing text generation...") + + test_cases = [ + { + "prompt": "Hello, how are you?", + "max_tokens": 50, + "temperature": 0.7 + }, + { + "prompt": "Tell me a joke", + "max_tokens": 100, + "temperature": 0.9 + }, + { + "prompt": "What is artificial intelligence?", + "max_tokens": 150, + "temperature": 0.5 + } + ] + + for i, test_case in enumerate(test_cases, 1): + print(f"\nTest {i}: {test_case['prompt']}") + + try: + response = requests.post( + "http://localhost:8000/generate", + json={ + "prompts": [test_case["prompt"]], + "max_tokens": test_case["max_tokens"], + "temperature": test_case["temperature"] + }, + timeout=30 + ) + + if response.status_code == 200: + result = response.json() + text = result["outputs"][0]["text"] + print(f"โœ… Response: {text}") + else: + print(f"โŒ Error: {response.status_code} - {response.text}") + + except requests.exceptions.RequestException as e: + print(f"โŒ Request failed: {e}") + + +def cleanup(): + """Clean up Docker container.""" + print("\n๐Ÿงน Cleaning up...") + try: + subprocess.run(['docker', 'stop', 'vllm-omni-cpu'], + capture_output=True, text=True) + subprocess.run(['docker', 'rm', 'vllm-omni-cpu'], + capture_output=True, text=True) + print("โœ… Cleanup completed") + except Exception as e: + print(f"โš ๏ธ Cleanup warning: {e}") + + +def main(): + """Main example function.""" + print("๐Ÿณ vLLM-omni Docker Example") + print("=" * 40) + + # Check Docker + if not check_docker_running(): + print("โŒ Docker is not running. Please start Docker Desktop.") + sys.exit(1) + + print("โœ… Docker is running") + + # Check if image exists + result = subprocess.run(['docker', 'images', 'vllm-omni-cpu:latest'], + capture_output=True, text=True) + if 'vllm-omni-cpu' not in result.stdout: + print("โŒ Docker image not found. Please build it first:") + print(" docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu .") + sys.exit(1) + + print("โœ… Docker image found") + + try: + # Start server + if not start_docker_server(): + sys.exit(1) + + # Test generation + test_generation() + + print("\n๐ŸŽ‰ Docker example completed successfully!") + + except KeyboardInterrupt: + print("\nโน๏ธ Example interrupted by user") + finally: + # Ask if user wants to keep server running + keep_running = input("\n๐Ÿค” Keep the server running? (y/N): ").lower().strip() + if keep_running != 'y': + cleanup() + else: + print("โœ… Server kept running. Access it at http://localhost:8000") + + +if __name__ == "__main__": + main() diff --git a/examples/basic/simple_usage.py b/examples/basic/simple_usage.py new file mode 100644 index 000000000..573548180 --- /dev/null +++ b/examples/basic/simple_usage.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +""" +Simple vLLM-omni usage example. + +This example demonstrates how to use vLLM-omni for basic text generation +with a single AR (Autoregressive) stage. +""" + +import asyncio +from vllm_omni.core.omni_llm import OmniLLM, AsyncOmniLLM +from vllm_omni.config import create_ar_stage_config + + +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"] + ) + + # Initialize OmniLLM + omni_llm = OmniLLM([stage_config]) + + # Prepare stage arguments + 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: + for completion in output.outputs: + print(f" Text: {completion.text}") + print(f" Finished: {completion.finish_reason != 'length'}") + print(f" Tokens: {len(completion.token_ids)} tokens") + + +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"] + ) + + # Initialize AsyncOmniLLM + omni_llm = AsyncOmniLLM([stage_config]) + + # Prepare stage arguments + 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: + for completion in output.outputs: + print(f" Text: {completion.text}") + print(f" Finished: {completion.finish_reason != 'length'}") + print(f" Tokens: {len(completion.token_ids)} tokens") + + +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"] + ) + + # Initialize OmniLLM + omni_llm = OmniLLM([stage_config]) + + # Multiple prompts + prompts = [ + "Tell me a joke", + "What's the weather like?", + "Explain quantum computing" + ] + + for prompt in prompts: + print(f"\nInput: {prompt}") + 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: + completion = outputs[0].outputs[0] + print(f"Output: {completion.text}") + + +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/scripts/docker-setup-macos.sh b/scripts/docker-setup-macos.sh new file mode 100755 index 000000000..aeb4fb662 --- /dev/null +++ b/scripts/docker-setup-macos.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +# vLLM-omni Docker Setup Script for macOS +# This script helps set up Docker for vLLM-omni on macOS + +set -e + +echo "๐Ÿณ vLLM-omni Docker Setup for macOS" +echo "==================================" + +# Check if Docker is installed +if ! command -v docker &> /dev/null; then + echo "โŒ Docker is not installed. Please install Docker Desktop for Mac:" + echo " https://www.docker.com/products/docker-desktop/" + exit 1 +fi + +# Check if Docker is running +if ! docker info &> /dev/null; then + echo "โŒ Docker is not running. Please start Docker Desktop." + exit 1 +fi + +echo "โœ… Docker is installed and running" + +# Check if NVIDIA Docker runtime is available (optional for macOS) +if docker info | grep -q nvidia; then + echo "โœ… NVIDIA Docker runtime detected" + GPU_SUPPORT=true +else + echo "โš ๏ธ NVIDIA Docker runtime not detected (normal on macOS)" + GPU_SUPPORT=false +fi + +# Create necessary directories +echo "๐Ÿ“ Creating directories..." +mkdir -p models logs notebooks + +# Build the Docker image +echo "๐Ÿ”จ Building vLLM-omni Docker image..." +docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu:latest . + +if [ $? -eq 0 ]; then + echo "โœ… Docker image built successfully" +else + echo "โŒ Failed to build Docker image" + exit 1 +fi + +# Test the image +echo "๐Ÿงช Testing Docker image..." +docker run --rm vllm-omni-cpu:latest python -c "import vllm_omni; print('vLLM-omni imported successfully')" + +if [ $? -eq 0 ]; then + echo "โœ… Docker image test passed" +else + echo "โŒ Docker image test failed" + exit 1 +fi + +echo "" +echo "๐ŸŽ‰ Setup complete! You can now use:" +echo "" +echo " # Start the server:" +echo " docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu" +echo "" +echo " # Start with Jupyter:" +echo " docker-compose -f docker/docker-compose.cpu.yml --profile jupyter up" +echo "" +echo " # Run a single command:" +echo " docker run --rm -p 8000:8000 vllm-omni-cpu:latest vllm serve Qwen/Qwen3-0.6B --omni --port 8000" +echo "" +echo " # Interactive shell:" +echo " docker run -it --rm vllm-omni-cpu:latest bash" +echo "" + +# Optional: Start the server +read -p "๐Ÿš€ Start the vLLM-omni server now? (y/N): " -n 1 -r +echo +if [[ $REPLY =~ ^[Yy]$ ]]; then + echo "Starting vLLM-omni server..." + docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu +fi diff --git a/scripts/download_models.py b/scripts/download_models.py new file mode 100755 index 000000000..382e6f7d3 --- /dev/null +++ b/scripts/download_models.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +""" +Model download script for vLLM-omni +Downloads commonly used AR and DiT models +""" + +import os +import sys +import argparse +from pathlib import Path + +def download_ar_models(models=None): + """Download AR (Autoregressive) models""" + if models is None: + models = [ + "Qwen/Qwen3-0.6B", + ] + + print("๐Ÿš€ Downloading AR models...") + + try: + from transformers import AutoTokenizer, AutoModel + except ImportError: + print("โŒ transformers not installed. Installing...") + os.system("pip install transformers") + from transformers import AutoTokenizer, AutoModel + + for model_name in models: + print(f"\n๐Ÿ“ฅ Downloading AR model: {model_name}") + try: + print(f" - Downloading tokenizer...") + tokenizer = AutoTokenizer.from_pretrained(model_name) + + print(f" - Downloading model...") + model = AutoModel.from_pretrained(model_name) + + print(f"โœ… {model_name} downloaded successfully") + + # Print model info + total_params = sum(p.numel() for p in model.parameters()) + size_mb = total_params * 4 / (1024 * 1024) + print(f" ๐Ÿ“Š Parameters: {total_params:,} ({size_mb:.1f} MB)") + + except Exception as e: + print(f"โŒ Failed to download {model_name}: {e}") + +def download_dit_models(models=None): + """Download DiT (Diffusion Transformer) models""" + if models is None: + models = [ + "stabilityai/stable-diffusion-2-1", + "runwayml/stable-diffusion-v1-5" + ] + + print("\n๐Ÿš€ Downloading DiT models...") + + try: + from diffusers import StableDiffusionPipeline + except ImportError: + print("โŒ diffusers not installed. Installing...") + os.system("pip install diffusers") + from diffusers import StableDiffusionPipeline + + for model_name in models: + print(f"\n๐Ÿ“ฅ Downloading DiT model: {model_name}") + try: + print(f" - Downloading pipeline...") + pipe = StableDiffusionPipeline.from_pretrained(model_name) + + print(f"โœ… {model_name} downloaded successfully") + + # Print model info + total_params = sum(p.numel() for p in pipe.unet.parameters()) + size_mb = total_params * 4 / (1024 * 1024) + print(f" ๐Ÿ“Š UNet parameters: {total_params:,} ({size_mb:.1f} MB)") + + except Exception as e: + print(f"โŒ Failed to download {model_name}: {e}") + +def check_cache_directory(): + """Check and display cache directory info""" + cache_dir = Path.home() / ".cache" / "huggingface" / "hub" + print(f"๐Ÿ“ Cache directory: {cache_dir}") + + if cache_dir.exists(): + total_size = sum(f.stat().st_size for f in cache_dir.rglob('*') if f.is_file()) + size_gb = total_size / (1024**3) + print(f"๐Ÿ“Š Total cache size: {size_gb:.2f} GB") + + # List downloaded models + model_dirs = [d for d in cache_dir.iterdir() if d.is_dir() and d.name.startswith("models--")] + if model_dirs: + print(f"๐Ÿ“‹ Downloaded models ({len(model_dirs)}):") + for model_dir in sorted(model_dirs): + model_name = model_dir.name.replace("models--", "").replace("--", "/") + print(f" - {model_name}") + else: + print("๐Ÿ“ Cache directory does not exist yet") + +def clear_cache(): + """Clear model cache""" + cache_dir = Path.home() / ".cache" / "huggingface" / "hub" + + if cache_dir.exists(): + print(f"๐Ÿ—‘๏ธ Clearing cache directory: {cache_dir}") + import shutil + shutil.rmtree(cache_dir) + print("โœ… Cache cleared successfully") + else: + print("๐Ÿ“ Cache directory does not exist") + +def main(): + parser = argparse.ArgumentParser(description="Download models for vLLM-omni") + parser.add_argument("--ar-models", nargs="+", help="AR models to download") + parser.add_argument("--dit-models", nargs="+", help="DiT models to download") + parser.add_argument("--all", action="store_true", help="Download all default models") + parser.add_argument("--check-cache", action="store_true", help="Check cache directory") + parser.add_argument("--clear-cache", action="store_true", help="Clear cache directory") + + args = parser.parse_args() + + print("๐ŸŽฏ vLLM-omni Model Downloader") + print("=" * 40) + + if args.clear_cache: + clear_cache() + return + + if args.check_cache: + check_cache_directory() + return + + if args.all or args.ar_models: + download_ar_models(args.ar_models) + + if args.all or args.dit_models: + download_dit_models(args.dit_models) + + if not any([args.all, args.ar_models, args.dit_models, args.check_cache, args.clear_cache]): + print("โ“ No action specified. Use --help for options.") + print("\nQuick start:") + print(" python scripts/download_models.py --all") + print(" python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B") + print(" python scripts/download_models.py --check-cache") + +if __name__ == "__main__": + main() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index e04cd3b66..503d32417 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -93,7 +93,8 @@ def create_ar_stage_config( model_path: str, input_modalities: List[str] = None, output_modalities: List[str] = None, - vllm_config: Optional[VllmConfig] = None + vllm_config: Optional[VllmConfig] = None, + executor_class: Optional[Type[Executor]] = None ) -> OmniStageConfig: """Create a configuration for an AR (Autoregressive) stage.""" if input_modalities is None: @@ -101,13 +102,16 @@ def create_ar_stage_config( 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( stage_id=stage_id, engine_type="AR", model_path=model_path, input_modalities=input_modalities, output_modalities=output_modalities, - vllm_config=vllm_config + vllm_config=vllm_config, + executor_class=executor_class ) @@ -118,7 +122,8 @@ def create_dit_stage_config( output_modalities: List[str] = None, dit_config: Optional[DiTConfig] = None, cache_config: Optional[DiTCacheConfig] = None, - vllm_config: Optional[VllmConfig] = None + vllm_config: Optional[VllmConfig] = None, + executor_class: Optional[Type[Executor]] = None ) -> OmniStageConfig: """Create a configuration for a DiT (Diffusion Transformer) stage.""" if input_modalities is None: @@ -133,6 +138,8 @@ def create_dit_stage_config( num_inference_steps=50 ) + # For now, we'll create minimal configs + # vllm_config and executor_class will be handled by the LLM class return OmniStageConfig( stage_id=stage_id, engine_type="DiT", @@ -141,5 +148,6 @@ def create_dit_stage_config( output_modalities=output_modalities, vllm_config=vllm_config, dit_config=dit_config, - cache_config=cache_config + cache_config=cache_config, + executor_class=executor_class ) diff --git a/vllm_omni/core/omni_llm.py b/vllm_omni/core/omni_llm.py index 9b94ccf8a..f25f8d929 100644 --- a/vllm_omni/core/omni_llm.py +++ b/vllm_omni/core/omni_llm.py @@ -2,13 +2,12 @@ Core OmniLLM and AsyncOmniLLM classes for multi-stage processing. """ -import asyncio from typing import List, Dict, Any, Optional, Union, Callable from vllm.entrypoints.llm import LLM from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.llm_engine import LLMEngine -from vllm.outputs import RequestOutput -from vllm.outputs import LoRARequest +from vllm.outputs import RequestOutput, LoRARequest + from ..config import OmniStageConfig from .stage_manager import StageManager @@ -26,6 +25,20 @@ class OmniLLM(LLM): ): # Use the first stage's model as the default model for LLM default_model = stage_configs[0].model_path if stage_configs else "test-model" + + # Fix configuration validation issues + # Ensure max_num_batched_tokens is at least as large as max_model_len + if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs: + if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']: + kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] + elif 'max_model_len' in kwargs: + # If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len + kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] + else: + # Set reasonable defaults to avoid validation errors + kwargs['max_model_len'] = 2048 + kwargs['max_num_batched_tokens'] = 2048 + super().__init__(model=default_model, **kwargs) self.stage_configs = stage_configs self.log_stats = log_stats @@ -153,34 +166,54 @@ class OmniLLM(LLM): priority: Optional[List[int]] = None ) -> Any: """Execute a single stage.""" - # This is a simplified implementation - # In practice, this would involve proper request management - # and integration with vLLM's engine system + # Use the parent LLM's generate method for actual text generation + prompt = processed_input.get("prompt", "") + max_tokens = processed_input.get("max_tokens", 100) + temperature = processed_input.get("temperature", 0.7) - # For now, we'll return a mock output - # In the full implementation, this would call stage_engine.generate() - from vllm.outputs import RequestOutput, CompletionOutput + # Generate using the base LLM class + from vllm.sampling_params import SamplingParams - mock_output = CompletionOutput( - index=0, - text=processed_input.get("prompt", ""), - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="length" + sampling_params = SamplingParams( + max_tokens=max_tokens, + temperature=temperature, + top_p=processed_input.get("top_p", 1.0), + frequency_penalty=processed_input.get("frequency_penalty", 0.0), + presence_penalty=processed_input.get("presence_penalty", 0.0), + stop=processed_input.get("stop", None) ) - return RequestOutput( - request_id="mock_request", - prompt=processed_input.get("prompt", ""), - prompt_token_ids=[], - outputs=[mock_output], - finished=True - ) + # Use the parent class's generate method + outputs = super().generate([prompt], sampling_params) + + # Return the first output (we're processing one prompt at a time) + if outputs: + return outputs[0] + else: + # Fallback to mock output if generation fails + from vllm.outputs import RequestOutput, CompletionOutput + + mock_output = CompletionOutput( + index=0, + text="Generation failed", + token_ids=[], + cumulative_logprob=0.0, + logprobs=None, + finish_reason="error" + ) + + return RequestOutput( + request_id="fallback_request", + prompt=prompt, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[mock_output], + finished=True + ) -class AsyncOmniLLM(AsyncLLM): - """Extended AsyncLLM supporting multiple engines and stage-based processing.""" +class AsyncOmniLLM(LLM): + """Extended LLM class supporting multiple engines and stage-based processing.""" def __init__( self, @@ -188,14 +221,30 @@ class AsyncOmniLLM(AsyncLLM): log_stats: bool = False, **kwargs ): - # Use the first stage's model as the default model for AsyncLLM - default_model = stage_configs[0].model_path if stage_configs else "test-model" - super().__init__(model=default_model, **kwargs) + # Use the first stage's model for the base LLM + if stage_configs and stage_configs[0].model_path: + model = stage_configs[0].model_path + else: + model = "Qwen/Qwen3-0.6B" + + # Fix configuration validation issues + # Ensure max_num_batched_tokens is at least as large as max_model_len + if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs: + if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']: + kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] + elif 'max_model_len' in kwargs: + # If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len + kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] + else: + # Set reasonable defaults to avoid validation errors + kwargs['max_model_len'] = 2048 + kwargs['max_num_batched_tokens'] = 2048 + + super().__init__(model=model, **kwargs) self.stage_configs = stage_configs self.log_stats = log_stats self.stage_manager = StageManager(stage_configs, log_stats) self.output_processor = MultimodalOutputProcessor() - self._initialize_async_stage_engines() def _initialize_async_stage_engines(self) -> None: """Initialize AsyncLLM instances for each stage.""" @@ -308,25 +357,47 @@ class AsyncOmniLLM(AsyncLLM): priority: Optional[List[int]] = None ) -> Any: """Execute a single stage asynchronously.""" - # This is a simplified implementation - # In practice, this would involve proper async request management + # Use the parent LLM's generate method for actual text generation + prompt = processed_input.get("prompt", "") + max_tokens = processed_input.get("max_tokens", 100) + temperature = processed_input.get("temperature", 0.7) - # For now, we'll return a mock output - from vllm.outputs import RequestOutput, CompletionOutput + # Generate using the base LLM class + from vllm.sampling_params import SamplingParams - mock_output = CompletionOutput( - index=0, - text=processed_input.get("prompt", ""), - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="length" + sampling_params = SamplingParams( + max_tokens=max_tokens, + temperature=temperature, + top_p=processed_input.get("top_p", 1.0), + frequency_penalty=processed_input.get("frequency_penalty", 0.0), + presence_penalty=processed_input.get("presence_penalty", 0.0), + stop=processed_input.get("stop", None) ) - return RequestOutput( - request_id="mock_request", - prompt=processed_input.get("prompt", ""), - prompt_token_ids=[], - outputs=[mock_output], - finished=True - ) + # Use the parent class's generate method + outputs = super().generate([prompt], sampling_params) + + # Return the first output (we're processing one prompt at a time) + if outputs: + return outputs[0] + else: + # Fallback to mock output if generation fails + from vllm.outputs import RequestOutput, CompletionOutput + + mock_output = CompletionOutput( + index=0, + text="Generation failed", + token_ids=[], + cumulative_logprob=0.0, + logprobs=None, + finish_reason="error" + ) + + return RequestOutput( + request_id="fallback_request", + prompt=prompt, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[mock_output], + finished=True + ) diff --git a/vllm_omni/entrypoints/api_server.py b/vllm_omni/entrypoints/api_server.py index 66c5952a7..2a74047b6 100644 --- a/vllm_omni/entrypoints/api_server.py +++ b/vllm_omni/entrypoints/api_server.py @@ -53,19 +53,22 @@ async def generate(request: GenerateRequest): # Prepare stage arguments if request.stage_args is None: - # Create default stage arguments - stage_args = [] - for i, prompt in enumerate(request.prompts): - stage_args.append({ - "prompt": prompt, - "max_tokens": request.max_tokens, - "temperature": request.temperature - }) + # 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 + }] else: stage_args = request.stage_args # Generate using omni_llm - outputs = await omni_llm.generate_async(stage_args) + try: + outputs = await omni_llm.generate_async(stage_args) + except Exception as e: + # Fallback to synchronous generation + outputs = omni_llm.generate(stage_args) # Convert outputs to response format response_outputs = [] diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 4a4753f14..bb2447f3c 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -12,8 +12,8 @@ def main(): """Main CLI entry point that intercepts vLLM commands.""" # Check if --omni flag is present if "--omni" in sys.argv: - # Remove --omni flag and process with vLLM-omni - omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"] + # Remove --omni flag and "serve" command, process with vLLM-omni + omni_args = [arg for arg in sys.argv[1:] if arg not in ["--omni", "serve"]] omni_serve = OmniServeCommand() omni_serve.run(omni_args) else: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index fab1f2d7d..5b6b17deb 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -21,10 +21,12 @@ class OmniServeCommand: description="vLLM-omni: Multi-modality models inference and serving" ) - # Model arguments + # Model arguments - make it optional with default parser.add_argument( "model", - help="Path to the model or model name" + nargs="?", + default="Qwen/Qwen3-0.6B", + help="Path to the model or model name (default: Qwen/Qwen3-0.6B)" ) # Server arguments @@ -105,19 +107,27 @@ class OmniServeCommand: stage_configs = [] stage_id = 0 - # Add AR stage if specified - if args.ar_stage: - ar_config = create_ar_stage_config( - stage_id=stage_id, - model_path=args.ar_stage, - input_modalities=["text"], - output_modalities=["text"] - ) - stage_configs.append(ar_config) - stage_id += 1 + # 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"] + ) + stage_configs.append(ar_config) + stage_id += 1 # Add DiT stage if specified if args.dit_stage: + dit_model = args.dit_stage + elif args.use_diffusers: + # Use a default DiT model if diffusers is enabled + dit_model = "stabilityai/stable-diffusion-2-1" + else: + dit_model = None + + if dit_model: dit_config = DiTConfig( model_type="dit", scheduler_type="ddpm", @@ -128,7 +138,7 @@ class OmniServeCommand: dit_stage_config = create_dit_stage_config( stage_id=stage_id, - model_path=args.dit_stage, + model_path=dit_model, input_modalities=["text"], output_modalities=["image"], dit_config=dit_config @@ -184,7 +194,7 @@ class OmniServeCommand: from vllm_omni.entrypoints.api_server import run_server await run_server( - omni_llm=omni_llm, + omni_llm_instance=omni_llm, host=args.host, port=args.port )