add precommit

.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

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

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:

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
    ```

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   
+For any questions or feedback regarding this design document, please contact  
-Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),   
+Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
 Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
 
 [image1]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAALYAAAFwCAYAAAAVNoEZAABxw0lEQVR4XuydBVRcWb6ve90709OStnTScel4x93dIUaEuBJ3dzciRPDgEtzdCRpCBCIEEkIgQHCCBHe+t6uSvNvNnZ55r2/fng5zvrX+69TZp4qkan3nz2/vKup88sm/OWqXL3HxwnnOnj7JubOnuaB6jqtXLqJ28TxXRWmrX8Ha0gIzE0OuX1bFQFsTcxMDLMyMMDXS44bWVRxsLLG5aYKuxlVsLMxwc3bAWF8bNzdnHGwt8fb2wNXZHlcnW3HMDjsbC4Ju3cLO2hwvT3ecnexwsLcmIMAHHy83+X18PVxxtLMUt+1xcbTBw8URRwdr+f1tLM1wd3PB08NZlAs24ue4uTrKf0bD5yfxb8qWDWvZtnE9W0Vt37yBvTu3cWDPTlE7OHviCBfPnsTEwECIfg7Vk0e5fkkVoxtamBneQF1NlSsXTmMrZJaN6elcF7LexFaIZ6SngbOjvfy2l5DUSwhoa2mK9U1jfL09ibwTKaS3IPx2BA4ONnh7uuHr4yUX1d/Pi0B/X1xd7PHz8cTT3UUc85SfIMHBQUJ4a3EfH/mJ4OsjThoXB/njPNydJbEl3rFisTKrli5i9fLFrF21jC0bVdi2aR07tqzn+KF9nD91HF1NDU4eO8CZIwe5fO4UBjoacpEvnDnG5fMnsDQ1Rk/rOje0rwmxLeQd3chAC0d7W6zMjfBwc8JHiGlmpIu5oS6hISHcibgtOrE99+4+wE50dR8hraeHm1zOW4G+BAb4yjtyYICf/FiAED0wMICw0DDxODvR8QNEBeLv7yO6txPeotN7e7lLYku8Y8akccycOhGlGVNRVprJsoXzWLlEmbUrl7Jv5xbOHD8sBD/ApnWrOLZvN2ePH0FDRJVronMf3rdDiH1aiK4p9s+JE0AWS6xE59bA2EAHO6ub3DTRw1UI7Ovjg/oVEWV01Llz567ouH6i+wbz8FEMtiKaBPr7yzuvn6+nED9I3rH9/bwJDgqUSy3r1OHh4YSEBOPu7io6/W1CQ0PFSeAvxHeTd3B/X19JbIl3TB4zjOnjRzJz8jiUFCazYPZ0li1QYrXo4nt3bObE4QPs3LIRlRWLObh7h9jfj5rqGdGtj7N/11Yh9hk0r14W+yfQ0bgiZLZC8/pljPR1ROwwwdxYF2d7GyGfN1cunkJfiB0h4oe3lweRkfeIjn4iF9vf1wcnkbX9fL0ICb4ll1kutrgtkzdEyB4ZGUlQ0C35Y+/cuSPv3oHimOxkkHX2W6KjN3x+Ev+mjBzQmzGD+zJxxCAmjxrC9AmjUBZyL1dWEkJvEF17G6tEB1++aB47Nq0Xsm+Vy71PbHdv2yTP4LL9Q3t3onVdTYhsyCXV02LyqCuyuS762texlk0oXV05e/IQetrqBInuK5v8RUc/JiLiDvYiisgmf7Y2N0UssZJv1a9fFQJ7yiUPkXXrsBAePHiAn+jMsghy//590dlDRRTxk8cW2VhoWLAktsQ7Bv7UleH9egi5+8kFl3Xw+TOnsmTBLDGpXMsOMaFcunCuiChKbF27ml1CbtnEcvt6FbnY504eE517uxB+CxpXLwlxNbhwXjbhvCGXWFOMmZsZibztwMmj+7khsniA6Kwurs4iSoi87OKCuuj4urqa3DQ3xdHBVt6pXZwd0NJQJ1AmbGiw6NC3xYkQLbK6lxA6WC55yP8VWxZrbhEWHiKJLfGOPl06MKhnVwb06Cy2XZg0aqiIJJNYNHcG61cvl2dtZXF78byZrF22iHUrl7B53WqRwxcI6dfLO/WW9Wvk3f3i2VOonj7BeRFT9HW0ULtwRoh/FIMb2hgZ6otMLuvql9klos3COYri3xrOiEH96dujK907d+Snrp3o37c3Gzdu4MKFCxjqG7Bn9y7uRt6Vd+gHUQ9xFsKHiMmnbEy29ReRRbaaEhYWJo8qDZ+fxL8pfbt2pF+3TqJ+ZHCvLiKODGX2tAly8VRWLJFPJBcoKcpFX6E8j9WL300slyrPkXd0WdfeIE4A2UqKbHlQtpJy7tQxtNWvckZILYspOprX0dFSZ8+2DezcqMKsCaOZOGwQvTp1oHPbVnRo2YL2rZrz/TdN+O6rJrRp8YMQvA+q5y9w9NARbgUFEf3wEffuR2FrZyUXOiIiQnTyEPmkUSa2LNLcu/dAElviHb06taF357bIBB/cqzNjh/YXOXs0cxWnsWjebPlKyVwRTebNmILyLAWWzJ3NMuW5LJijIO/o61YvY/WShWxUWcmB3dvZLgQ/emAvZ04cZc/2zULmTaiKuLJv22ZmjxvJqH496dK6pagW9Gnfhp7tWtH5h+9p//23tPruG1qI+v6br/n2qy/47K9/odn337FGRQUHJ2fCbt/GyFhPxI4Q+eRRJrivr7d8AvngQZTo6tGS2BLv6P6jELtLO7nY/bv/yPCBPZkiBJylMIV5QuS5M6czR3EKStMnMU8m+5yZLFGazdwZU1mzfAmrly2UTyzXr1omjyMb1qxgn8jch/fvYdsGFXZsXMMusV0i7j+6V1eGdm3PGBF5pvXtztxBPZjetyujxL8/VPw/erRuzo9C8pZNheBNv+b7r7/iy88/o2PHjixcspTbtyO4IbK4bMnvzp13KySyzB0cIptYRktiS/wXnzdpzpdNmtGtQyt+6tiK/j06MXbEYCaPG8WMKROYOWUc0yaOY6ro4orittKMacxWmMysaeNZtmAei+fPEZ19FquWLmSdkHv5ovnyzL1dxJQ1Qvr1ohaOG8rxJTO4ukYJ/a1LODprDEsGdmdkux8Y2KopHb/+knZffk6LJp/xvRD5e3G7+Zdf0PKbr2gtJG8jZP/uqy9Zu24tqmdOEn77jrxjy96wkb1BExERyd2794m4c1cSW+Idn33ZlL99IaJA6x/o1KaZ6N7tGTmkP+NGDpXLPG38KLnkk8eNYPqkMXKpZ0ydIKQfK2LJHHkHV1aaIYReIM/jMtFlgss691KRzeeKn6G7YwVup3fgr3YAw83KLO/VkbX9OrNhYBfmd23JrC6tmNT+B7p89zVdRafu1uwbOjf9ivbffUUHsS/r4h3F/69Fs2YsWbhAZOn78owdGBgol1vWyWWTyYg7EZLYEu9o/cVXrO71E6aL53Nx1mT6iLw9sE9Phg/qx/jhQ+Q1doTYjhryXu6xTJNtx4scrjCVOSKiKImoIsvispJ19MULlEQGn8GsYf05vWI2+aGOZHqb4LJvNY+v7iP2+hEeXdpL2NF1uG9V5ti4nuwb1pUlPVoz88cWTGzblHGihrVuypA23zOsfXOGCPE7t/qBZt9+g7qmFmFh4XKxZXlb9kbPbdG1ZZ87afj8JP5NGd2mFfuHD8BCeQ43F8/hzMxxzB3al5H9fmJo/z4MH9CXEYMHMEpMKscNFxFl9AjGjxjERNHRFSbIJB/N9MljmDV9MjOnTUJh8nhmT5/CxNHDMNu/luI77pRGB5Dhb0t1zB2qnoRREuxEnpsx989sI3T/apzXKXBz6WRMl01Bf9F4VKcP5uj4Pmwd1p21g7ui3Lsj83p1oEerZrQV3fv7pt9y4uRxbgUHERIWRlh4BMFC9FvBoZLYEu/YObgfp0YO4ZLowtemjhViKXB1wVT2KIxBccRAxorOPVgIPnRQX8YMHSjv4GPeSz5p7EgmjhkhjywyoadPkmXxMfJMPmHoAOJstKhLjKLiaTiVCQ+pSY6jJj6K0jA3CjxNSdI/S4L2MR5c2sW9CzsIOqKC156lGC2bivr8MRye0I+dI3uyelBXlvXvxMA2zWnb7Ds+//QvjBkziiAxaQwJDRWTx3D5kqC/9Ja6xAcU27Rkcgvx67/59yzo2JrDo/thvVYZj51rCT20leCDm7DfvALN5XNRHj2UcQP7MKRPV9HJe4suLuvkQvZRw+SCTxg9XH57pDgZji6bSXVcJHWZidS+SaMuN43alGdysasfBlER6UNpsCNv/azJctQh1eoacdrHuX9hF67bFmK+SpFz04dwZHw/to3sxfqh3ZgssniXZl/LlwT/+p//ia+vn5g8+uPt7SX/YJSrq4sktsQ7un39Fd2/+ZoB337NpJZNWdezI/pKY3FSmU+o2mVCjc3xd/LCw9kHbSNrDl/QQFFpAaNHDGNov3dZfIzo3rLJ5tiR7/L4zJGDCBQ5ui4tntr8TOpF1WWnUJsUQ+2z+1RHBVF1L4DSUFeKAuzJdjEkzUaT+BtniLq8F+ftypitVOCsEPvoxH7sHt2XzUN7MLtHGwa0aEpLIfYXf/sMS0sLPDw8cHR0xN7ORl4Nn5/Evymd2rXgxzY/0EF0wj5dO3NE2xTj5CLMC+rQL6zDtKQe+/I6rEvrMS6owfBNFQY5leimlnLI0pOl2/aLDN6X0UPede9TG1dxS/sMiO5cHO5DcYgHed4OZNkZURbqSfWDYKrvBlN+25/iIDcK/Rx542JEpp0uTzWPEXl+Jw6b52O8fDrnFYZxcupg9o7pK7p2Dxb3aYdil5a0+eYrvvnicw4e3I/lTQsM9G+gr6eL3g1dSWyJd3T44Rt+bN2C+bsOc9w7ApMiuFkMN/JrsSyuwbm8FuuiWkzF/s3Caizf1mCZX4VpbgV6yW+5/iCRpVv3MkFMJBdMn0BehDf50SGQLaLH6xfUy+JHVCiVYZ6UeltRGexO9T2xHxFI0S2RtX0cyHMxI9P2Bo9Flw8/vQ3bjfMwWDqVc0LsE1MGsXd0H7aP6M6SPu2Z3bUVbb9pwjdffsnKlcvQ1rjO1SuXuHZVVpclsf8gNomSvdgfKuaXh//1qD3O4fLrKjRya9HJrUE3uxqr4mrcy2vwLKvDMKsKSyG0s9i3Lxad+20txqJj3xRlJeQ2yyhBNy6Dq0GPWLN8Kbn3g6jNEdEjL5uHphZcO3QRNyNHEnyCqY2NpfbFC2rjnlBxL0x0clkMuUmSmQZPdS9wR3U/t4Line truncated

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:

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?",
+        "What is the capital of France?", max_tokens=30, temperature=0.3
-        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 = [
+    prompts = ["Tell me a joke", "What's 2+2?", "Describe the color blue"]
-        "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/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 = [{
+    stage_args = [
-        "prompt": "Hello, how are you today?",
+        {
-        "max_tokens": 50,
+            "prompt": "Hello, how are you today?",
-        "temperature": 0.7
+            "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 = [{
+    stage_args = [
-        "prompt": "What is artificial intelligence?",
+        {
-        "max_tokens": 100,
+            "prompt": "What is artificial intelligence?",
-        "temperature": 0.8
+            "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 = [{
+        stage_args = [{"prompt": prompt, "max_tokens": 30, "temperature": 0.9}]
-            "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!")

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

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

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
-

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
 }

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

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",
+            name="test_tensor", shape=[1, 512, 512], dtype="float32", persistent=True
-            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
+        enable_optimization=True,
-    )
+    )

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(
+        dit_config = DiTConfig(num_inference_steps=50)
-            model_type="dit",
-            scheduler_type="ddpm",
-            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(
+        config = DiTConfig(num_inference_steps=50)
-            model_type="dit",
-            scheduler_type="ddpm",
-            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(
+        config = create_ar_stage_config(stage_id=0, model_path="test-model")
-            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(
+        dit_config = DiTConfig(num_inference_steps=50)
-            model_type="dit",
+
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
-        
         config = create_dit_stage_config(
-            stage_id=1,
+            stage_id=1, model_path="test-dit-model", dit_config=dit_config
-            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
-

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
+This package extends vLLM beyond traditional text-based, autoregressive
-to support multi-modality models with non-autoregressive structures and non-textual outputs.
+generation to support multi-modality models with non-autoregressive
+structures and non-textual outputs.
 
 Architecture:
-- 🟡 Modified: vLLM components modified for multimodal support  
+- 🟡 Modified: vLLM components modified for multimodal support
-- 🔴 Added: New components for multimodal and non-autoregressive processing
+- 🔴 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.
-]
+]

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

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)
 

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
+This module provides caching functionality for DiT (Diffusion Transformer)
-to optimize inference performance and memory usage.
+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(), 
+        oldest_key = min(
-                        key=lambda k: self.cache_tensors[k].timestamp)
+            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(),
+        lru_key = min(
-                     key=lambda k: self.cache_tensors[k].access_count)
+            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(),
+        lfu_key = min(
-                     key=lambda k: self.cache_tensors[k].access_count)
+            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
+        self.cache_timeout = timeout

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
+    This scheduler extends vLLM's scheduler to support multimodal and
-    processing with additional fields and methods specific to vLLM-omni.
+    non-autoregressive processing with additional fields and methods
+    specific to vLLM-omni.
     """
 
     # Ensure scheduled_new_reqs carry omni-specific payloads (e.g., additional_information)

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 (
+from ..worker.gpu_diffusion_worker import DiffusionGPUWorker, DiffusionRunnerOutput
-    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 = (
+                config_model_path = model_cfg.get("model") or model_cfg.get(
-                    model_cfg.get("model")
+                    "model_path"
-                    or model_cfg.get("model_path")
                 )
             else:
-                config_model_path = (
+                config_model_path = getattr(model_cfg, "model", None) or getattr(
-                    getattr(model_cfg, "model", None)
+                    model_cfg, "model_path", 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)
+            float(guidance_scale) if guidance_scale is not None else self.guidance_scale
-            if guidance_scale is not None
-            else self.guidance_scale
         )
 
         return self.worker.generate(

vllm_omni/entrypoints/api_server.py

@@ -3,17 +3,18 @@ API server for vLLM-omni.
 """
 
 import asyncio
-from typing import Dict, Any, List
+from typing import Any, Dict, List
-from fastapi import FastAPI, HTTPException
+
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
 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 = [{
+            stage_args = [
-                "prompt": " ".join(request.prompts) if request.prompts else "",
+                {
-                "max_tokens": request.max_tokens,
+                    "prompt": " ".join(request.prompts) if request.prompts else "",
-                "temperature": request.temperature
+                    "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({
+                    response_outputs.append(
-                        "text": getattr(out, 'text', ''),
+                        {
-                        "finished": getattr(out, 'finish_reason', 'length') != 'length',
+                            "text": getattr(out, "text", ""),
-                        "tokens": getattr(out, 'token_ids', [])
+                            "finished": getattr(out, "finish_reason", "length")
-                    })
+                            != "length",
+                            "tokens": getattr(out, "token_ids", []),
+                        }
+                    )
             else:
-                response_outputs.append({
+                response_outputs.append({"text": "", "finished": True, "tokens": []})
-                    "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(
+    config = uvicorn.Config(app=app, host=host, port=port, log_level="info")
-        app=app,
+
-        host=host,
-        port=port,
-        log_level="info"
-    )
-    
     server = uvicorn.Server(config)
     await server.serve()
 

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

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(
+        parser.add_argument("--ar-stage", type=str, help="AR stage model path")
-            "--ar-stage",
+
-            type=str,
+        parser.add_argument("--dit-stage", type=str, help="DiT stage model path")
-            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",
+            "--log-stats", action="store_true", help="Enable logging statistics"
-            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,
+            stage_configs=stage_configs, log_stats=parsed_args.log_stats
-            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(
+            await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port)
-                omni_llm_instance=omni_llm,
-                host=args.host,
-                port=args.port
-            )
         except KeyboardInterrupt:
             print("\nShutting down server...")
         except Exception as e: