feat: enhance serving functionality and add comprehensive examples

- Improve API server with better error handling and response formatting
- Enhance CLI with additional options for DiT stages and configuration
- Add comprehensive examples in examples/basic/ including:
  - API client with health checks and text generation
  - Docker setup and usage examples
  - Simple usage patterns for different scenarios
- Add utility scripts for model downloading and Docker setup
- Update documentation with implementation details and testing guidelines
- Fix configuration validation issues in OmniLLM
- Improve stage configuration handling for AR and DiT stages
- Add proper error handling and fallback mechanisms

Tested with Qwen3-0.6B model:
- Server starts successfully on port 8000
- Health and info endpoints working correctly
- Text generation with various parameters functioning
- API client examples working as expected
- CLI help and configuration options working properly

.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

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).

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"]
 )

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"]
             ),

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

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
+

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()
+

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()

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!")

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

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()

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
     )

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
+            )

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 = []

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:

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
             )