Merge pull request #32 from Gaohan123/precommit_set

add precommit and solve conflict

.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

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -16,7 +16,7 @@ body:
       - Passwords or authentication credentials
       - Private URLs or endpoints
       - Personal or confidential data
-      
+
       Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
 - type: textarea
   attributes:
@@ -58,7 +58,7 @@ body:
           input_modalities=["text"],
           output_modalities=["text"]
       )
-      
+
       dit_config = create_dit_stage_config(
           stage_id=1,
           model_path="stabilityai/stable-diffusion-2-1",

.gitignore

@@ -176,7 +176,6 @@ tests/
 
 # vLLM-omni specific
 # Model files and checkpoints
-models/
 checkpoints/
 *.bin
 *.safetensors

.pre-commit-config.yaml

@@ -0,0 +1,29 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.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.13.2
+    hooks:
+      - id: isort
+
+  - repo: https://github.com/psf/black
+    rev: 24.10.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.1.1
+    hooks:
+      - id: flake8
+        args: ["--max-line-length=88", "--extend-ignore=E203,W503"]

CONTRIBUTING.md

@@ -0,0 +1,19 @@
+# Contributing to vLLM
+
+## Getting Started
+
+Setup environment with uv
+```bash
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+Pre-commit has been linked to git commit. For run checks locally, please install pre-commit:
+```bash
+uv pip install pre-commit
+```
+
+Then run:
+```bash
+pre-commit run --show-diff-on-failure --color=always --all-files
+```

README.md

@@ -72,4 +72,4 @@ The output audio is saved in ./output_audio
 - [ ] Online inference support
 - [ ] Support for other models
 
-For detailed model management, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).
+For detailed model management, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).

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   
-Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),   
+For any questions or feedback regarding this design document, please contact  
+Han Gao: [hgaoaf@connect.ust.hk](mailto:hgaoaf@connect.ust.hk),  
 Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
 
 [image1]: <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?",
-        max_tokens=30,
-        temperature=0.3
+        "What is the capital of France?", max_tokens=30, temperature=0.3
     )
     if result:
         text = result["outputs"][0]["text"]
-        print(f"   Input: What is the capital of France?")
+        print("   Input: What is the capital of France?")
         print(f"   Output: {text}")
-    
+
     # Example 4: Multiple prompts
     print("\n4. Multiple prompts:")
-    prompts = [
-        "Tell me a joke",
-        "What's 2+2?",
-        "Describe the color blue"
-    ]
-    
+    prompts = ["Tell me a joke", "What's 2+2?", "Describe the color blue"]
+
     for prompt in prompts:
         result = generate_text(prompt, max_tokens=30)
         if result:
             text = result["outputs"][0]["text"]
             print(f"   Q: {prompt}")
             print(f"   A: {text}")
-    
+
     # Example 5: Show full response structure
     print("\n5. Full response structure:")
     result = generate_text("Hello world", max_tokens=20)
     if result:
         print("   Full API response:")
         print(json.dumps(result, indent=2))
-    
+
     print("\n" + "=" * 40)
     print("API client examples completed!")
 
 
 if __name__ == "__main__":
     main()
-

examples/basic/text_generation.py

@@ -7,40 +7,43 @@ with a single AR (Autoregressive) stage.
 """
 
 import asyncio
-from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
+
 from vllm_omni.config import create_ar_stage_config
+from vllm_omni.entrypoints.omni_llm import AsyncOmniLLM, OmniLLM
 
 
 def sync_example():
     """Synchronous usage example."""
     print("=== Synchronous vLLM-omni Example ===")
-    
+
     # Create a simple AR stage configuration
     stage_config = create_ar_stage_config(
         stage_id=0,
         model_path="Qwen/Qwen3-0.6B",
         input_modalities=["text"],
-        output_modalities=["text"]
+        output_modalities=["text"],
     )
-    
+
     # Initialize OmniLLM
     omni_llm = OmniLLM([stage_config])
-    
+
     # Prepare stage arguments
-    stage_args = [{
-        "prompt": "Hello, how are you today?",
-        "max_tokens": 50,
-        "temperature": 0.7
-    }]
-    
+    stage_args = [
+        {
+            "prompt": "Hello, how are you today?",
+            "max_tokens": 50,
+            "temperature": 0.7,
+        }
+    ]
+
     # Generate text
     print(f"Input: {stage_args[0]['prompt']}")
     outputs = omni_llm.generate(stage_args)
-    
+
     # Display results
     for i, output in enumerate(outputs):
         print(f"Output {i}:")
-        if hasattr(output, 'outputs') and output.outputs:
+        if hasattr(output, "outputs") and output.outputs:
             for completion in output.outputs:
                 print(f"  Text: {completion.text}")
                 print(f"  Finished: {completion.finish_reason != 'length'}")
@@ -50,33 +53,35 @@ def sync_example():
 async def async_example():
     """Asynchronous usage example."""
     print("\n=== Asynchronous vLLM-omni Example ===")
-    
+
     # Create a simple AR stage configuration
     stage_config = create_ar_stage_config(
         stage_id=0,
         model_path="Qwen/Qwen3-0.6B",
         input_modalities=["text"],
-        output_modalities=["text"]
+        output_modalities=["text"],
     )
-    
+
     # Initialize AsyncOmniLLM
     omni_llm = AsyncOmniLLM([stage_config])
-    
+
     # Prepare stage arguments
-    stage_args = [{
-        "prompt": "What is artificial intelligence?",
-        "max_tokens": 100,
-        "temperature": 0.8
-    }]
-    
+    stage_args = [
+        {
+            "prompt": "What is artificial intelligence?",
+            "max_tokens": 100,
+            "temperature": 0.8,
+        }
+    ]
+
     # Generate text asynchronously
     print(f"Input: {stage_args[0]['prompt']}")
     outputs = await omni_llm.generate_async(stage_args)
-    
+
     # Display results
     for i, output in enumerate(outputs):
         print(f"Output {i}:")
-        if hasattr(output, 'outputs') and output.outputs:
+        if hasattr(output, "outputs") and output.outputs:
             for completion in output.outputs:
                 print(f"  Text: {completion.text}")
                 print(f"  Finished: {completion.finish_reason != 'length'}")
@@ -86,36 +91,32 @@ async def async_example():
 def multi_prompt_example():
     """Example with multiple prompts."""
     print("\n=== Multi-Prompt Example ===")
-    
+
     # Create stage configuration
     stage_config = create_ar_stage_config(
         stage_id=0,
         model_path="Qwen/Qwen3-0.6B",
         input_modalities=["text"],
-        output_modalities=["text"]
+        output_modalities=["text"],
     )
-    
+
     # Initialize OmniLLM
     omni_llm = OmniLLM([stage_config])
-    
+
     # Multiple prompts
     prompts = [
         "Tell me a joke",
         "What's the weather like?",
-        "Explain quantum computing"
+        "Explain quantum computing",
     ]
-    
+
     for prompt in prompts:
         print(f"\nInput: {prompt}")
-        stage_args = [{
-            "prompt": prompt,
-            "max_tokens": 30,
-            "temperature": 0.9
-        }]
-        
+        stage_args = [{"prompt": prompt, "max_tokens": 30, "temperature": 0.9}]
+
         outputs = omni_llm.generate(stage_args)
-        
-        if outputs and hasattr(outputs[0], 'outputs') and outputs[0].outputs:
+
+        if outputs and hasattr(outputs[0], "outputs") and outputs[0].outputs:
             completion = outputs[0].outputs[0]
             print(f"Output: {completion.text}")
 
@@ -123,15 +124,15 @@ def multi_prompt_example():
 if __name__ == "__main__":
     print("vLLM-omni Simple Usage Examples")
     print("=" * 40)
-    
+
     # Run synchronous example
     sync_example()
-    
+
     # Run asynchronous example
     asyncio.run(async_example())
-    
+
     # Run multi-prompt example
     multi_prompt_example()
-    
+
     print("\n" + "=" * 40)
     print("Examples completed!")

examples/offline_inference/qwen_2_5_omni/README.md

@@ -34,4 +34,4 @@ The output audio is saved in ./output_audio
 - [ ] Adaptation from current vllm branch to stable vllm v0.11.0
 - [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests
 - [ ] Online inference support
-- [ ] Support for other models
+- [ ] Support for other models

examples/offline_inference/qwen_2_5_omni/end2end.py

@@ -1,18 +1,17 @@
 import argparse
 import os
-import soundfile as sf
+import os as _os_env_toggle
 import random
-import numpy as np
-import torch
 
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
 from vllm.sampling_params import SamplingParams
 
-import os as _os_env_toggle
-_os_env_toggle.environ["VLLM_USE_V1"] = "1"
-
 from vllm_omni.entrypoints.omni_llm import OmniLLM
-from utils import make_omni_prompt
 
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
 
 SEED = 42
 # Set all random seeds
@@ -33,39 +32,84 @@ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', required=True, help='Path to merged model directory (will be created if downloading).')
-    parser.add_argument('--thinker-model', type=str, default=None)
-    parser.add_argument('--talker-model', type=str, default=None)
-    parser.add_argument('--code2wav-model', type=str, default=None)
-    parser.add_argument('--hf-hub-id', default='Qwen/Qwen2.5-Omni-7B', help='Hugging Face repo id to download if needed.')
-    parser.add_argument('--hf-revision', default=None, help='Optional HF revision (branch/tag/commit).')
-    parser.add_argument('--prompts', required=True, nargs='+', help='Input text prompts.')
-    parser.add_argument('--voice-type', default='default', help='Voice type, e.g., m02, f030, default.')
-    parser.add_argument('--code2wav-dir', default=None, help='Path to code2wav folder (contains spk_dict.pt).')
-    parser.add_argument('--dit-ckpt', default=None, help='Path to DiT checkpoint file (e.g., dit.pt).')
-    parser.add_argument('--bigvgan-ckpt', default=None, help='Path to BigVGAN checkpoint file.')
-    parser.add_argument('--dtype', default='bfloat16', choices=['float16', 'bfloat16', 'float32'])
-    parser.add_argument('--max-model-len', type=int, default=32768)
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument(
+        "--hf-revision", default=None, help="Optional HF revision (branch/tag/commit)."
+    )
+    parser.add_argument(
+        "--prompts", required=True, nargs="+", help="Input text prompts."
+    )
+    parser.add_argument(
+        "--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
+    )
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument(
+        "--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt)."
+    )
+    parser.add_argument(
+        "--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file."
+    )
+    parser.add_argument(
+        "--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"]
+    )
+    parser.add_argument("--max-model-len", type=int, default=32768)
 
     parser.add_argument("--thinker-only", action="store_true")
     parser.add_argument("--text-only", action="store_true")
     parser.add_argument("--do-wave", action="store_true")
-    parser.add_argument('--prompt_type',
-                        choices=[
-                            'text', 'audio', 'audio-long', 'audio-long-chunks',
-                            'audio-long-expand-chunks', 'image', 'video',
-                            'video-frames', 'audio-in-video', 'audio-in-video-v2',
-                            "audio-multi-round", "badcase-vl", "badcase-text",
-                            "badcase-image-early-stop", "badcase-two-audios",
-                            "badcase-two-videos", "badcase-multi-round",
-                            "badcase-voice-type", "badcase-voice-type-v2",
-                            "badcase-audio-tower-1", "badcase-audio-only"
-                        ],
-                        default='text')
-    parser.add_argument('--use-torchvision', action='store_true')
-    parser.add_argument('--tokenize', action='store_true')
-    parser.add_argument('--output-wav', default="output.wav", help='Output wav file path.')
-    parser.add_argument('--thinker-hidden-states-dir', default="thinker_hidden_states", help='Path to thinker hidden states directory.')
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav", default="output.wav", help="Output wav file path."
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
     args = parser.parse_args()
     return args
 
@@ -75,38 +119,40 @@ def main():
     model_name = args.model
     omni_llm = OmniLLM(model=model_name)
     thinker_sampling_params = SamplingParams(
-                                            temperature=0.0,    # Deterministic - no randomness
-                                            top_p=1.0,          # Disable nucleus sampling
-                                            top_k=-1,           # Disable top-k sampling
-                                            max_tokens=2048,
-                                            seed=SEED,          # Fixed seed for sampling
-                                            detokenize=True,
-                                            repetition_penalty=1.1,
-                                            )
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
     talker_sampling_params = SamplingParams(
-                                            temperature=0.0,    # Deterministic - no randomness
-                                            top_p=1.0,          # Disable nucleus sampling
-                                            top_k=-1,           # Disable top-k sampling
-                                            max_tokens=2048,
-                                            seed=SEED,          # Fixed seed for sampling
-                                            detokenize=True,
-                                            repetition_penalty=1.1,
-                                            stop_token_ids=[8294]
-                                            )
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+        stop_token_ids=[8294],
+    )
     code2wav_sampling_params = SamplingParams(
-                                            temperature=0.0,    # Deterministic - no randomness
-                                            top_p=1.0,          # Disable nucleus sampling
-                                            top_k=-1,           # Disable top-k sampling
-                                            max_tokens=2048,
-                                            seed=SEED,          # Fixed seed for sampling
-                                            detokenize=True,
-                                            repetition_penalty=1.1,
-                                            )
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
 
-    sampling_params_list = [thinker_sampling_params,
-                            talker_sampling_params,
-                            code2wav_sampling_params]
-    
     prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
     omni_outputs = omni_llm.generate(prompt, sampling_params_list)
 
@@ -121,10 +167,14 @@ def main():
             for output in stage_outputs.request_output:
                 request_id = output.request_id
                 audio_tensor = output.multimodal_output["audio"]
-                output_wav = os.path.join(args.output_wav, f"output_{output.request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                output_wav = os.path.join(
+                    args.output_wav, f"output_{output.request_id}.wav"
+                )
+                sf.write(
+                    output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
+                )
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/offline_inference/qwen_2_5_omni/processing_omni.py

@@ -8,6 +8,7 @@ import time
 import warnings
 from functools import lru_cache
 from io import BytesIO
+
 import requests
 import torch
 import torchvision
@@ -16,7 +17,6 @@ from PIL import Image
 from torchvision import io, transforms
 from torchvision.transforms import InterpolationMode
 
-
 logger = logging.getLogger(__name__)
 
 IMAGE_FACTOR = 28
@@ -43,20 +43,24 @@ def round_by_factor(number: int, factor: int) -> int:
 
 
 def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
     return math.ceil(number / factor) * factor
 
 
 def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
     return math.floor(number / factor) * factor
 
 
-def smart_resize(height: int,
-                 width: int,
-                 factor: int = IMAGE_FACTOR,
-                 min_pixels: int = MIN_PIXELS,
-                 max_pixels: int = MAX_PIXELS) -> tuple[int, int]:
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
     """
     Rescales the image so that the following conditions are met:
 
@@ -68,7 +72,8 @@ def smart_resize(height: int,
     """
     if max(height, width) / min(height, width) > MAX_RATIO:
         raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
+            f"got {max(height, width) / min(height, width)}"
         )
     h_bar = max(factor, round_by_factor(height, factor))
     w_bar = max(factor, round_by_factor(width, factor))
@@ -83,8 +88,9 @@ def smart_resize(height: int,
     return h_bar, w_bar
 
 
-def fetch_image(ele: dict[str, str | Image.Image],
-                size_factor: int = IMAGE_FACTOR) -> Image.Image:
+def fetch_image(
+    ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR
+) -> Image.Image:
     if "image" in ele:
         image = ele["image"]
     else:
@@ -105,10 +111,11 @@ def fetch_image(ele: dict[str, str | Image.Image],
         image_obj = Image.open(image)
     if image_obj is None:
         raise ValueError(
-            f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
+            f"Unrecognized image input, support local path, http url, "
+            f"base64 and PIL.Image, got {image}"
         )
     image = image_obj.convert("RGB")
-    ## resize
+    # resize
     if "resized_height" in ele and "resized_width" in ele:
         resized_height, resized_width = smart_resize(
             ele["resized_height"],
@@ -143,8 +150,10 @@ def smart_nframes(
             support either `fps` or `nframes`:
                 - nframes: the number of frames to extract for model inputs.
                 - fps: the fps to extract frames for model inputs.
-                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
-                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
         total_frames (int): the original total number of frames of the video.
         video_fps (int | float): the original fps of the video.
 
@@ -154,34 +163,39 @@ def smart_nframes(
     Returns:
         int: the number of frames for video used for model inputs.
     """
-    assert not ("fps" in ele
-                and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    assert not (
+        "fps" in ele and "nframes" in ele
+    ), "Only accept either `fps` or `nframes`"
     if "nframes" in ele:
         nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
     else:
         fps = ele.get("fps", FPS)
-        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES),
-                                    FRAME_FACTOR)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
         max_frames = floor_by_factor(
             ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
-            FRAME_FACTOR)
+            FRAME_FACTOR,
+        )
         nframes = total_frames / video_fps * fps
         nframes = min(max(nframes, min_frames), max_frames)
         nframes = round_by_factor(nframes, FRAME_FACTOR)
     if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
         raise ValueError(
-            f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}."
+            f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], "
+            f"but got {nframes}."
         )
     return nframes
 
 
-def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
     """read video using torchvision.io.read_video
 
     Args:
         ele (dict): a dict contains the configuration of video.
         support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
             - video_start: the start time of video.
             - video_end: the end time of video.
     Returns:
@@ -191,7 +205,8 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
     if version.parse(torchvision.__version__) < version.parse("0.19.0"):
         if "http://" in video_path or "https://" in video_path:
             warnings.warn(
-                "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0."
+                "torchvision < 0.19.0 does not support http/https video path, "
+                "please upgrade to 0.19.0."
             )
         if "file://" in video_path:
             video_path = video_path[7:]
@@ -206,11 +221,10 @@ def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
     total_frames, video_fps = video.size(0), info["video_fps"]
     total_duration = round(total_frames / video_fps, 3)
     logger.info(
-        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, "
+        f"duration={total_duration}s, time={time.time() - st:.3f}s"
     )
-    nframes = smart_nframes(ele,
-                            total_frames=total_frames,
-                            video_fps=video_fps)
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long()
     video = video[idx]
     return video, total_duration, nframes
@@ -222,34 +236,38 @@ def is_decord_available() -> bool:
     return importlib.util.find_spec("decord") is not None
 
 
-def _read_video_decord(ele: dict, ) -> torch.Tensor:
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
     """read video using decord.VideoReader
 
     Args:
         ele (dict): a dict contains the configuration of video.
         support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
             - video_start: the start time of video.
             - video_end: the end time of video.
     Returns:
         torch.Tensor: the video tensor with shape (T, C, H, W).
     """
     import decord
+
     video_path = ele["video"]
     st = time.time()
     vr = decord.VideoReader(video_path)
     # TODO: support start_pts and end_pts
-    if 'video_start' in ele or 'video_end' in ele:
+    if "video_start" in ele or "video_end" in ele:
         raise NotImplementedError(
-            "not support start_pts and end_pts in decord for now.")
+            "not support start_pts and end_pts in decord for now."
+        )
     total_frames, video_fps = len(vr), vr.get_avg_fps()
     total_duration = round(total_frames / video_fps, 3)
     logger.info(
-        f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s"
+        f"decord:  {video_path=}, {total_frames=}, {video_fps=}, "
+        f"time={time.time() - st:.3f}s"
     )
-    nframes = smart_nframes(ele,
-                            total_frames=total_frames,
-                            video_fps=video_fps)
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
     idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
     video = vr.get_batch(idx).asnumpy()
     video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
@@ -272,27 +290,27 @@ def get_video_reader_backend() -> str:
         video_reader_backend = "decord"
     else:
         video_reader_backend = "torchvision"
-    # print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
     return video_reader_backend
 
 
 def fetch_video(
-        ele: dict,
-        image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+    ele: dict, image_factor: int = IMAGE_FACTOR
+) -> torch.Tensor | list[Image.Image]:
     if isinstance(ele["video"], str):
         video_reader_backend = get_video_reader_backend()
-        video, total_dur, nframes = VIDEO_READER_BACKENDS[
-            video_reader_backend](ele)
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
         frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
         grid_timestamps = frame_timestamps[::FRAME_FACTOR]
         second_per_grid = grid_timestamps[1] - grid_timestamps[0]
         nframes, _, height, width = video.shape
-        factor = spatial_patch_size * spatial_merge_size
         min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
         total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
         max_pixels = max(
             min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
-            int(min_pixels * 1.05))
+            int(min_pixels * 1.05),
+        )
         max_pixels = ele.get("max_pixels", max_pixels)
         # min_pixels = (factor ** 2) * 52
         # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
@@ -323,11 +341,9 @@ def fetch_video(
         process_info.pop("type", None)
         process_info.pop("video", None)
         images = [
-            fetch_image({
-                "image": video_element,
-                **process_info
-            },
-                        size_factor=image_factor)
+            fetch_image(
+                {"image": video_element, **process_info}, size_factor=image_factor
+            )
             for video_element in ele["video"]
         ]
         nframes = ceil_by_factor(len(images), FRAME_FACTOR)
@@ -336,8 +352,7 @@ def fetch_video(
         return images, None, None, None
 
 
-def extract_vision_info(
-        conversations: list[dict] | list[list[dict]]) -> list[dict]:
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
     vision_infos = []
     if isinstance(conversations[0], dict):
         conversations = [conversations]
@@ -345,18 +360,21 @@ def extract_vision_info(
         for message in conversation:
             if isinstance(message["content"], list):
                 for ele in message["content"]:
-                    if ("image" in ele or "image_url" in ele or "video" in ele
-                            or ele["type"] in ("image", "image_url", "video")):
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
                         vision_infos.append(ele)
     return vision_infos
 
 
 def process_vision_info(
     conversations: list[dict] | list[list[dict]],
-) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]]
-           | None]:
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
     vision_infos = extract_vision_info(conversations)
-    ## Read images or videos
+    # Read images or videos
     image_inputs = []
     video_inputs = []
     for vision_info in vision_infos:
@@ -370,4 +388,4 @@ def process_vision_info(
         image_inputs = None
     if len(video_inputs) == 0:
         video_inputs = None
-    return image_inputs, video_inputs
+    return image_inputs, video_inputs

examples/offline_inference/qwen_2_5_omni/run.sh

@@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \
                                  --dit-ckpt none \
                                  --bigvgan-ckpt none \
                                  --output-wav output_audio \
-                                 --prompt_type text
+                                 --prompt_type text

examples/offline_inference/qwen_2_5_omni/utils.py

@@ -1,30 +1,33 @@
 import tempfile
+from typing import Dict, List, Optional, Union
 from urllib.request import urlopen
-import librosa
-import soundfile as sf
-import resampy
-from typing import Dict, Optional
-import torch
-import requests
-import torchvision.io
 
-from typing import Union, List
-from vllm.inputs import TextPrompt
-from vllm_omni.inputs.data import OmniTokensPrompt
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
 from processing_omni import fetch_image, fetch_video
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
 
 
 def get_system_prompt():
-    
+
     return {
-        'role':
-        'system',
-        'content': [{
-            'type':
-            'text',
-            'text':
-            'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.'
-        }]
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
     }
 
 
@@ -34,9 +37,7 @@ def resample_wav_to_16khz(input_filepath):
     if len(data.shape) > 1:
         data = data[:, 0]
     # resample to 16kHz
-    data_resampled = resampy.resample(data,
-                                      sr_orig=original_sample_rate,
-                                      sr_new=16000)
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
     return data_resampled
 
 
@@ -55,8 +56,7 @@ def fetch_and_read_video(args, video_url: str, fps=2):
         total_duration = round(total_frames / video_fps, 3)
         nframes = int(total_frames / video_fps * fps)
 
-        frame_timestamps = total_duration * torch.arange(1,
-                                                         nframes + 1) / nframes
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
         grid_timestamps = frame_timestamps[::2]
         second_per_grid = grid_timestamps[1] - grid_timestamps[0]
 
@@ -70,7 +70,8 @@ def fetch_and_read_video(args, video_url: str, fps=2):
 
     def read_video_with_transformers(video_file_name: Union[str, List[str]]):
         video, total_duration, nframes, second_per_grid = fetch_video(
-            {'video': video_file_name})
+            {"video": video_file_name}
+        )
         if total_duration is None and nframes is None:
             nframes = len(video)
             total_duration = 0.5 * nframes
@@ -89,7 +90,10 @@ def fetch_and_read_video(args, video_url: str, fps=2):
     if isinstance(video_url, str) and video_url.startswith("http"):
         with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
             resp = requests.get(video_url)
-            assert resp.status_code == requests.codes.ok, f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
+            assert resp.status_code == requests.codes.ok, (
+                f"Failed to fetch video from {video_url}, "
+                f"status_code:{resp.status_code}, resp:{resp}"
+            )
 
             temp_video_file.write(urlopen(video_url).read())
             temp_video_file_path = temp_video_file.name
@@ -106,104 +110,104 @@ def make_inputs_qwen2_omni(
     use_audio_in_video: Optional[bool] = False,
     tokenize: bool = False,
 ) -> Union[OmniTokensPrompt, TextPrompt]:
-    
-    from transformers import AutoConfig, AutoProcessor, AutoTokenizer
+
+    from transformers import AutoConfig, AutoProcessor
+
     processor = AutoProcessor.from_pretrained(args.model)
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
 
     try:
         config = AutoConfig.from_pretrained(args.model)
-        if 'Qwen2_5OmniModel' in config.architectures:
+        if "Qwen2_5OmniModel" in config.architectures:
             args.legacy_omni_video = False
         else:
             args.legacy_omni_video = True
-    except:
+    except Exception:
         args.legacy_omni_video = True
 
     audios, images, videos = [], [], []
     for message in messages:
-        if not isinstance(message['content'], list):
-            message['content'] = [{
-                'type': 'text',
-                'text': message['content'],
-            }]
-        index, num_contents = 0, len(message['content'])
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
         while index < num_contents:
-            ele = message['content'][index]
-            if 'type' not in ele:
-                if 'text' in ele:
-                    ele['type'] = 'text'
-                elif 'audio' in ele:
-                    ele['type'] = 'audio'
-                elif 'audio_url' in ele:
-                    ele['type'] = 'audio_url'
-                elif 'image' in ele:
-                    ele['type'] = 'image'
-                elif 'image_url' in ele:
-                    ele['type'] = 'image_url'
-                elif 'video' in ele:
-                    ele['type'] = 'video'
-                elif 'video_url' in ele:
-                    ele['type'] = 'video_url'
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
                 else:
-                    raise ValueError(f'Unknown ele: {ele}')
+                    raise ValueError(f"Unknown ele: {ele}")
 
-            if ele['type'] == 'audio' or ele['type'] == 'audio_url':
-                if 'audio_url' in ele:
-                    audio_key = 'audio_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_audio_file:
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
                         temp_audio_file.write(urlopen(ele[audio_key]).read())
                         temp_audio_file_path = temp_audio_file.name
-                        audios.append(
-                            resample_wav_to_16khz(temp_audio_file_path))
-                        ele['audio'] = temp_audio_file_path
-                elif 'audio' in ele:
-                    audio_key = 'audio'
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
                     audios.append(resample_wav_to_16khz(ele[audio_key]))
                 else:
-                    raise ValueError(f'Unknown ele {ele}')
-            elif use_audio_in_video and (ele['type'] == 'video'
-                                         or ele['type'] == 'video_url'):
+                    raise ValueError(f"Unknown ele {ele}")
+            elif use_audio_in_video and (
+                ele["type"] == "video" or ele["type"] == "video_url"
+            ):
                 # use video as audio as well
-                if 'video_url' in ele:
-                    audio_key = 'video_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_video_file:
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
                         temp_video_file.write(urlopen(ele[audio_key]).read())
                         temp_video_file_path = temp_video_file.name
                         ele[audio_key] = temp_video_file_path
-                        audios.append(
-                            librosa.load(temp_video_file_path, sr=16000)[0])
-                        videos.append(
-                            fetch_and_read_video(args, temp_video_file_path))
-                        ele['video'] = temp_video_file_path
-                elif 'video' in ele:
-                    audio_key = 'video'
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
                     audios.append(librosa.load(ele[audio_key], sr=16000)[0])
                     videos.append(fetch_and_read_video(args, audio_key))
                 else:
                     raise ValueError("Unknown ele {}".format(ele))
                 # insert a audio after the video
-                message['content'].insert(index + 1, {
-                    "type": "audio",
-                    "audio": ele[audio_key],
-                })
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
                 # no need to load the added audio again
                 index += 1
-            elif ele['type'] == 'video' or ele['type'] == 'video_url':
-                if 'video_url' in ele:
-                    video_key = 'video_url'
-                    with tempfile.NamedTemporaryFile(
-                            delete=True) as temp_video_file:
-                        temp_video_file.write(urlopen(ele['video_url']).read())
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
                         temp_video_file_path = temp_video_file.name
                         videos.append(fetch_and_read_video(args, temp_video_file))
-                        ele['video'] = temp_video_file_path
+                        ele["video"] = temp_video_file_path
                 else:
-                    video_key = 'video'
+                    video_key = "video"
                     videos.append(fetch_and_read_video(args, ele[video_key]))
-            elif ele['type'] == 'image' or ele['type'] == 'image_url':
+            elif ele["type"] == "image" or ele["type"] == "image_url":
                 images.append(fetch_image(ele))
 
             # move to the next content
@@ -249,11 +253,8 @@ def make_text_prompt(args, prompt):
         {
             "role": "user",
             "content": [
-                {
-                    "type": "text",
-                    "text": prompt
-                },
-            ]
+                {"type": "text", "text": prompt},
+            ],
         },
     ]
 
@@ -264,26 +265,29 @@ def make_text_prompt(args, prompt):
 def make_audio_in_video_v2_prompt(args):
     messages = [
         {
-            'role':
-            'system',
-            'content': [{
-                'type':
-                'text',
-                'text':
-                'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.'
-            }]
-        },
-        {
-            "role":
-            "user",
+            "role": "system",
             "content": [
                 {
-                    "type":
-                    "video_url",
-                    "video_url":
-                    "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": (
+                        "https://qianwen-res.oss-cn-beijing.aliyuncs.com/"
+                        "Qwen2.5-Omni/draw_small.mp4"
+                    ),
                 },
-            ]
+            ],
         },
     ]
     prompt = make_inputs_qwen2_omni(
@@ -295,11 +299,13 @@ def make_audio_in_video_v2_prompt(args):
     return prompt
 
 
-def make_omni_prompt(args, prompt = None) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
-    if args.prompt_type == 'text':
+def make_omni_prompt(
+    args, prompt=None
+) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
         prompt = make_text_prompt(args, prompt)
-    elif args.prompt_type == 'audio-in-video-v2':
+    elif args.prompt_type == "audio-in-video-v2":
         prompt = make_audio_in_video_v2_prompt(args)
     else:
-        raise ValueError(f'Unsupported prompt type: {args.prompt_type}')
-    return prompt
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt

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",
-            shape=[1, 512, 512],
-            dtype="float32",
-            persistent=True
+            name="test_tensor", shape=[1, 512, 512], dtype="float32", persistent=True
         )
     ]
-    
+
     return DiTCacheConfig(
         cache_tensors=cache_tensors,
         max_cache_size=1024 * 1024 * 1024,  # 1GB
         cache_strategy="fifo",
-        enable_optimization=True
-    )
+        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(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
+        dit_config = DiTConfig(num_inference_steps=50)
         config = OmniStageConfig(
             stage_id=1,
             engine_type="DiT",
             model_path="test-dit-model",
             input_modalities=["text"],
             output_modalities=["image"],
-            dit_config=dit_config
+            dit_config=dit_config,
         )
         assert config.engine_type == "DiT"
         assert config.dit_config is not None
-        assert config.dit_config.model_type == "dit"
-    
+        assert config.dit_config.num_inference_steps == 50
+
     def test_invalid_engine_type(self):
         """Test validation of engine type."""
         with pytest.raises(ValueError):
@@ -56,9 +53,9 @@ class TestOmniStageConfig:
                 engine_type="INVALID",
                 model_path="test-model",
                 input_modalities=["text"],
-                output_modalities=["text"]
+                output_modalities=["text"],
             )
-    
+
     def test_empty_modalities(self):
         """Test validation of empty modalities."""
         with pytest.raises(ValueError):
@@ -67,18 +64,18 @@ class TestOmniStageConfig:
                 engine_type="AR",
                 model_path="test-model",
                 input_modalities=[],
-                output_modalities=["text"]
+                output_modalities=["text"],
             )
-        
+
         with pytest.raises(ValueError):
             OmniStageConfig(
                 stage_id=0,
                 engine_type="AR",
                 model_path="test-model",
                 input_modalities=["text"],
-                output_modalities=[]
+                output_modalities=[],
             )
-    
+
     def test_dit_requires_config(self):
         """Test that DiT engine requires dit_config."""
         with pytest.raises(ValueError):
@@ -87,32 +84,26 @@ class TestOmniStageConfig:
                 engine_type="DiT",
                 model_path="test-dit-model",
                 input_modalities=["text"],
-                output_modalities=["image"]
+                output_modalities=["image"],
             )
 
 
 class TestDiTConfig:
     def test_dit_config_defaults(self):
         """Test DiT configuration with defaults."""
-        config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
+        config = DiTConfig(num_inference_steps=50)
         assert config.use_diffusers is False
         assert config.diffusers_pipeline is None
         assert config.guidance_scale == 7.5
         assert config.height == 512
         assert config.width == 512
-    
+
     def test_diffusers_config(self):
         """Test DiT configuration with diffusers."""
         config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
             num_inference_steps=50,
             use_diffusers=True,
-            diffusers_pipeline="stable-diffusion"
+            diffusers_pipeline="stable-diffusion",
         )
         assert config.use_diffusers is True
         assert config.diffusers_pipeline == "stable-diffusion"
@@ -126,16 +117,16 @@ class TestDiTCacheConfig:
                 name="test_tensor",
                 shape=[1, 512, 512],
                 dtype="float32",
-                persistent=True
+                persistent=True,
             )
         ]
-        
+
         config = DiTCacheConfig(
             cache_tensors=cache_tensors,
             max_cache_size=1024 * 1024 * 1024,
-            cache_strategy="fifo"
+            cache_strategy="fifo",
         )
-        
+
         assert len(config.cache_tensors) == 1
         assert config.max_cache_size == 1024 * 1024 * 1024
         assert config.cache_strategy == "fifo"
@@ -145,35 +136,25 @@ class TestDiTCacheConfig:
 class TestHelperFunctions:
     def test_create_ar_stage_config(self):
         """Test create_ar_stage_config helper function."""
-        config = create_ar_stage_config(
-            stage_id=0,
-            model_path="test-model"
-        )
-        
+        config = create_ar_stage_config(stage_id=0, model_path="test-model")
+
         assert config.stage_id == 0
         assert config.engine_type == "AR"
         assert config.model_path == "test-model"
         assert config.input_modalities == ["text"]
         assert config.output_modalities == ["text"]
-    
+
     def test_create_dit_stage_config(self):
         """Test create_dit_stage_config helper function."""
-        dit_config = DiTConfig(
-            model_type="dit",
-            scheduler_type="ddpm",
-            num_inference_steps=50
-        )
-        
+        dit_config = DiTConfig(num_inference_steps=50)
+
         config = create_dit_stage_config(
-            stage_id=1,
-            model_path="test-dit-model",
-            dit_config=dit_config
+            stage_id=1, model_path="test-dit-model", dit_config=dit_config
         )
-        
+
         assert config.stage_id == 1
         assert config.engine_type == "DiT"
         assert config.model_path == "test-dit-model"
         assert config.input_modalities == ["text"]
         assert config.output_modalities == ["image"]
         assert config.dit_config is not None
-

vllm_omni/__init__.py

@@ -1,45 +1,46 @@
 """
-vLLM-omni: Multi-modality models inference and serving with non-autoregressive structures.
+vLLM-omni: Multi-modality models inference and serving with
+non-autoregressive structures.
 
-This package extends vLLM beyond traditional text-based, autoregressive generation
-to support multi-modality models with non-autoregressive structures and non-textual outputs.
+This package extends vLLM beyond traditional text-based, autoregressive
+generation to support multi-modality models with non-autoregressive
+structures and non-textual outputs.
 
 Architecture:
-- 🟡 Modified: vLLM components modified for multimodal support  
-- 🔴 Added: New components for multimodal and non-autoregressive processing
+- 🟡 Modified: vLLM components modified for multimodal support
+- 🔴 Added: New components for multimodal and non-autoregressive
+  processing
 """
 
 __version__ = "0.1.0"
 __author__ = "vLLM-omni Team"
 __email__ = "hsliuustc@gmail.com"
 
-# Main entry points
-from .entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
 from .config import (
-    OmniStageConfig,
-    DiTConfig,
     DiTCacheConfig,
+    DiTConfig,
+    OmniStageConfig,
     create_ar_stage_config,
     create_dit_stage_config,
 )
 
+# Main entry points
+from .entrypoints.omni_llm import AsyncOmniLLM, OmniLLM
+
 __all__ = [
     # Version info
     "__version__",
-    "__author__", 
+    "__author__",
     "__email__",
-    
     # Main components
     "OmniLLM",
     "AsyncOmniLLM",
-    
     # Configuration
     "OmniStageConfig",
     "DiTConfig",
     "DiTCacheConfig",
     "create_ar_stage_config",
     "create_dit_stage_config",
-    
     # All other components are available through their respective modules
     # processors.*, schedulers.*, executors.*, etc.
-]
+]

vllm_omni/config/__init__.py

@@ -1,20 +1,20 @@
 """
 Configuration module for vLLM-omni.
 """
-from typing import Optional
-from pydantic.dataclasses import dataclass
-from pydantic import ConfigDict
 
-from vllm.config import ModelConfig
-from vllm.config import config
+from typing import Optional
+
+from pydantic import ConfigDict
+from pydantic.dataclasses import dataclass
+from vllm.config import ModelConfig, 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,
 )
@@ -24,16 +24,16 @@ from .stage_config import (
 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
 class OmniModelConfig(ModelConfig):
     """Configuration for Omni models, extending the base ModelConfig."""
-    
+
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
-    
+
     @property
     def registry(self):
         return me_models.OmniModelRegistry
-    
+
     @property
     def architectures(self) -> list[str]:
         return [self.model_arch]
@@ -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
-to optimize inference performance and memory usage.
+This module provides caching functionality for DiT (Diffusion Transformer)
+models to optimize inference performance and memory usage.
 """
 
 import time
+from typing import Any, Dict, List, Optional
+
 import torch
-from typing import Dict, List, Optional, Any
-from dataclasses import dataclass
+
 from ..config import DiTCacheConfig, DiTCacheTensor
 
 
 class DiTCacheManager:
     """Manages DiT-specific caching for optimized inference."""
-    
+
     def __init__(self, config: DiTCacheConfig):
         self.config = config
         self.cache_tensors: Dict[str, DiTCacheTensor] = {}
@@ -24,30 +25,30 @@ class DiTCacheManager:
         self.cache_hits = 0
         self.cache_misses = 0
         self.cache_timeout = 3600  # 1 hour default timeout
-    
+
     def allocate_cache(self, request_id: str, size: int) -> torch.Tensor:
         """Allocate cache for a specific request."""
         # Check if we have enough space
         if self.current_cache_size + size > self.max_cache_size:
             self._evict_cache()
-        
+
         # Create tensor
         tensor = torch.zeros(size, dtype=torch.float32)
-        
+
         # Store in cache
         cache_tensor = DiTCacheTensor(
             name=request_id,
             tensor=tensor,
             timestamp=time.time(),
             access_count=0,
-            size_bytes=size * 4  # Assuming float32
+            size_bytes=size * 4,  # Assuming float32
         )
-        
+
         self.cache_tensors[request_id] = cache_tensor
         self.current_cache_size += cache_tensor.size_bytes
-        
+
         return tensor
-    
+
     def get_cache(self, request_id: str) -> Optional[torch.Tensor]:
         """Get cached tensor for a request."""
         if request_id in self.cache_tensors:
@@ -58,58 +59,58 @@ class DiTCacheManager:
         else:
             self.cache_misses += 1
             return None
-    
+
     def store_cache(self, request_id: str, tensor: torch.Tensor) -> None:
         """Store a tensor in the cache."""
         if request_id in self.cache_tensors:
             # Update existing cache
             old_tensor = self.cache_tensors[request_id]
             self.current_cache_size -= old_tensor.size_bytes
-            
+
             new_size_bytes = tensor.numel() * 4  # Assuming float32
             self.current_cache_size += new_size_bytes
-            
+
             self.cache_tensors[request_id] = DiTCacheTensor(
                 name=request_id,
                 tensor=tensor,
                 timestamp=time.time(),
                 access_count=old_tensor.access_count,
-                size_bytes=new_size_bytes
+                size_bytes=new_size_bytes,
             )
         else:
             # Create new cache entry
             new_size_bytes = tensor.numel() * 4
             if self.current_cache_size + new_size_bytes > self.max_cache_size:
                 self._evict_cache()
-            
+
             self.cache_tensors[request_id] = DiTCacheTensor(
                 name=request_id,
                 tensor=tensor,
                 timestamp=time.time(),
                 access_count=0,
-                size_bytes=new_size_bytes
+                size_bytes=new_size_bytes,
             )
             self.current_cache_size += new_size_bytes
-    
+
     def release_cache(self, request_id: str) -> None:
         """Release cache for a request."""
         if request_id in self.cache_tensors:
             cache_tensor = self.cache_tensors[request_id]
             self.current_cache_size -= cache_tensor.size_bytes
             del self.cache_tensors[request_id]
-    
+
     def clear_expired_cache(self) -> None:
         """Clear expired cache entries."""
         current_time = time.time()
         expired_keys = []
-        
+
         for request_id, cache_tensor in self.cache_tensors.items():
             if current_time - cache_tensor.timestamp > self.cache_timeout:
                 expired_keys.append(request_id)
-        
+
         for request_id in expired_keys:
             self.release_cache(request_id)
-    
+
     def _evict_cache(self) -> None:
         """Evict cache entries based on strategy."""
         if self.config.cache_strategy == "fifo":
@@ -120,42 +121,48 @@ class DiTCacheManager:
             self._evict_lfu()
         else:
             self._evict_fifo()  # Default to FIFO
-    
+
     def _evict_fifo(self) -> None:
         """Evict cache entries using FIFO strategy."""
         if not self.cache_tensors:
             return
-        
+
         # Find oldest entry
-        oldest_key = min(self.cache_tensors.keys(), 
-                        key=lambda k: self.cache_tensors[k].timestamp)
+        oldest_key = min(
+            self.cache_tensors.keys(),
+            key=lambda k: self.cache_tensors[k].timestamp,
+        )
         self.release_cache(oldest_key)
-    
+
     def _evict_lru(self) -> None:
         """Evict cache entries using LRU strategy."""
         if not self.cache_tensors:
             return
-        
+
         # Find least recently used entry
-        lru_key = min(self.cache_tensors.keys(),
-                     key=lambda k: self.cache_tensors[k].access_count)
+        lru_key = min(
+            self.cache_tensors.keys(),
+            key=lambda k: self.cache_tensors[k].access_count,
+        )
         self.release_cache(lru_key)
-    
+
     def _evict_lfu(self) -> None:
         """Evict cache entries using LFU strategy."""
         if not self.cache_tensors:
             return
-        
+
         # Find least frequently used entry
-        lfu_key = min(self.cache_tensors.keys(),
-                     key=lambda k: self.cache_tensors[k].access_count)
+        lfu_key = min(
+            self.cache_tensors.keys(),
+            key=lambda k: self.cache_tensors[k].access_count,
+        )
         self.release_cache(lfu_key)
-    
+
     def get_statistics(self) -> Dict[str, Any]:
         """Get cache statistics."""
         total_requests = self.cache_hits + self.cache_misses
         hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0
-        
+
         return {
             "cache_size": self.current_cache_size,
             "max_cache_size": self.max_cache_size,
@@ -163,16 +170,16 @@ class DiTCacheManager:
             "cache_hits": self.cache_hits,
             "cache_misses": self.cache_misses,
             "hit_rate": hit_rate,
-            "num_cached_tensors": len(self.cache_tensors)
+            "num_cached_tensors": len(self.cache_tensors),
         }
-    
+
     def clear_all_cache(self) -> None:
         """Clear all cache entries."""
         self.cache_tensors.clear()
         self.current_cache_size = 0
         self.cache_hits = 0
         self.cache_misses = 0
-    
+
     def set_cache_timeout(self, timeout: float) -> None:
         """Set cache timeout in seconds."""
-        self.cache_timeout = timeout
+        self.cache_timeout = timeout

vllm_omni/core/sched/__init__.py

@@ -2,13 +2,12 @@
 Scheduling components for vLLM-omni.
 """
 
-from .scheduler import OmniScheduler
 from .diffusion_scheduler import DiffusionScheduler
 from .output import OmniNewRequestData
+from .scheduler import OmniScheduler
 
 __all__ = [
     "OmniScheduler",
     "DiffusionScheduler",
     "OmniNewRequestData",
 ]
-

vllm_omni/core/sched/diffusion_scheduler.py

@@ -1,21 +1,30 @@
-from vllm.v1.core.sched.scheduler import SchedulerOutput, EngineCoreOutputs, Request, RequestStatus, SpecDecodingStats
+import time
 from collections import defaultdict
 from typing import Optional
-from vllm.v1.core.sched.request_queue import create_request_queue
-from vllm.v1.engine import EngineCoreEventType
+
 from vllm.distributed.kv_events import KVEventBatch
-import time
+from vllm.v1.core.sched.request_queue import create_request_queue
+from vllm.v1.core.sched.scheduler import (
+    EngineCoreOutputs,
+    Request,
+    RequestStatus,
+    SchedulerOutput,
+    SpecDecodingStats,
+)
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+
+from vllm_omni.core.sched.output import OmniNewRequestData
 from vllm_omni.core.sched.scheduler import OmniScheduler
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm.v1.engine import EngineCoreOutput
-from vllm_omni.core.sched.output import OmniNewRequestData
 
 
 class DiffusionScheduler(OmniScheduler):
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
-        - Feed all input tokens of the request at once (if 0, allocate 1 placeholder token).
-        - If the token budget cannot be satisfied at once, fall back to the default vLLM scheduling.
+        - Feed all input tokens of the request at once
+          (if 0, allocate 1 placeholder token).
+        - If the token budget cannot be satisfied at once, fall back to the
+          default vLLM scheduling.
         """
 
         # Select requests with zero prompt and using pooling
@@ -41,13 +50,15 @@ class DiffusionScheduler(OmniScheduler):
         # independent of pooling_params)
         while self.waiting and token_budget > 0 and capacity > 0:
             request = self.waiting.peek_request()
-            # Uniformly treat as diffusion. A feature flag can be added later via config or request tag.
+            # Uniformly treat as diffusion. A feature flag can be added later
+            # via config or request tag.
 
             # Allocate all input tokens for the request in one shot
             # (allocate 1 placeholder if zero)
             required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
             if required_tokens > token_budget:
-                # Insufficient budget to process all inputs at once; stop fast path attempt
+                # Insufficient budget to process all inputs at once;
+                # stop fast path attempt
                 break
             num_new_tokens = required_tokens
             new_blocks = self.kv_cache_manager.allocate_slots(
@@ -67,8 +78,7 @@ class DiffusionScheduler(OmniScheduler):
             self.running.append(request)
             request.status = RequestStatus.RUNNING
             if self.log_stats:
-                request.record_event(EngineCoreEventType.SCHEDULED,
-                                     scheduled_timestamp)
+                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
 
             req_to_new_block_ids[request.request_id] = new_blocks.get_block_ids()
             num_scheduled_tokens[request.request_id] = num_new_tokens
@@ -90,7 +100,9 @@ class DiffusionScheduler(OmniScheduler):
             any_request = self.running[0]
             num_common_prefix_blocks = (
                 self.kv_cache_manager.get_num_common_prefix_blocks(
-                    any_request, len(self.running)))
+                    any_request, len(self.running)
+                )
+            )
 
         grammar_bitmask = self.structured_output_manager.grammar_bitmask(
             self.requests,
@@ -137,16 +149,20 @@ class DiffusionScheduler(OmniScheduler):
             batch = KVEventBatch(ts=time.time(), events=events)
             self.kv_event_publisher.publish(batch)
 
-        # Update internal state (advance num_computed_tokens, free encoder inputs, etc.)
+        # Update internal state (advance num_computed_tokens, free encoder inputs,
+        # etc.)
         self._update_after_schedule(scheduler_output)
         return scheduler_output
+
     """
     Scheduler for the diffusion model.
     This scheduler is modified to stop the request immediately for the diffusion model.
     This is because the diffusion model can generate the final image/audio in one step.
-    Note: This is just a minimal modification to the original scheduler, and there should be some further efforts to optimize the scheduler. 
+    Note: This is just a minimal modification to the original scheduler,
+    and there should be some further efforts to optimize the scheduler.
     The original scheduler is still used for the AR model.
     """
+
     def update_from_output(
         self,
         scheduler_output: SchedulerOutput,
@@ -182,11 +198,13 @@ class DiffusionScheduler(OmniScheduler):
                 continue
 
             req_index = model_runner_output.req_id_to_index[req_id]
-            generated_token_ids = sampled_token_ids[
-                req_index] if sampled_token_ids else []
+            generated_token_ids = (
+                sampled_token_ids[req_index] if sampled_token_ids else []
+            )
 
             scheduled_spec_token_ids = (
-                scheduler_output.scheduled_spec_decode_tokens.get(req_id))
+                scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            )
             if scheduled_spec_token_ids:
                 # num_computed_tokens represents the number of tokens
                 # processed in the current step, considering scheduled
@@ -194,13 +212,15 @@ class DiffusionScheduler(OmniScheduler):
                 # num_computed_tokens is decreased by the number of rejected
                 # tokens, where is given by:
                 # len(scheduled_spec_token_ids) + 1 - len(generated_token_ids).
-                num_tokens_rejected = (len(scheduled_spec_token_ids) + 1 -
-                                       len(generated_token_ids))
+                num_tokens_rejected = (
+                    len(scheduled_spec_token_ids) + 1 - len(generated_token_ids)
+                )
                 request.num_computed_tokens -= num_tokens_rejected
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=len(scheduled_spec_token_ids),
-                    num_accepted_tokens=len(generated_token_ids) - 1)
+                    num_accepted_tokens=len(generated_token_ids) - 1,
+                )
 
             new_logprobs = None
             new_token_ids = generated_token_ids
@@ -212,7 +232,8 @@ class DiffusionScheduler(OmniScheduler):
 
             # Diffusion request: completes in one step; mark finished and free resources
             request.status = RequestStatus.FINISHED_STOPPED
-            # Optional: set a stop_reason for front-end clarity (does not affect protocol)
+            # Optional: set a stop_reason for front-end clarity
+            # (does not affect protocol)
             request.stop_reason = request.stop_reason or "diffusion_done"
             kv_transfer_params = self._free_request(request)
             if status_before_stop == RequestStatus.RUNNING:
@@ -221,19 +242,22 @@ class DiffusionScheduler(OmniScheduler):
                 stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None \
-                and request.sampling_params.logprobs is not None and logprobs:
+            if (
+                request.sampling_params is not None
+                and request.sampling_params.logprobs is not None
+                and logprobs
+            ):
                 # NOTE: once we support N tokens per step (spec decode),
                 # the outer lists can be of length > 1.
                 new_logprobs = logprobs.slice(req_index, req_index + 1)
 
-            if new_token_ids and self.structured_output_manager.should_advance(
-                    request):
-                # NOTE: structured_output_request
-                # should not be None if use_structured_output, we have
-                # check above, so safe to ignore type warning
-                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]
-                    req_id, new_token_ids)
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                # NOTE: structured_output_request should not be None if
+                # use_structured_output, we have check above, so safe to ignore
+                # type warning
+                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]  # noqa: E501
+                    req_id, new_token_ids
+                )
 
             # spec_token_ids comes from the model runner output
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -244,15 +268,15 @@ class DiffusionScheduler(OmniScheduler):
                 if self.structured_output_manager.should_advance(request):
                     metadata = request.structured_output_request
                     # Needs to happen after new_token_ids are accepted.
-                    request.spec_token_ids = metadata.grammar.validate_tokens(  # type: ignore[union-attr]
-                        spec_token_ids[req_index])
+                    request.spec_token_ids = metadata.grammar.validate_tokens(  # type: ignore[union-attr]  # noqa: E501
+                        spec_token_ids[req_index]
+                    )
                 else:
                     request.spec_token_ids = spec_token_ids[req_index]
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None \
-                or kv_transfer_params:
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
 
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
@@ -267,7 +291,8 @@ class DiffusionScheduler(OmniScheduler):
                         events=request.take_events(),
                         kv_transfer_params=kv_transfer_params,
                         num_cached_tokens=request.num_cached_tokens,
-                    ))
+                    )
+                )
 
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -284,8 +309,7 @@ class DiffusionScheduler(OmniScheduler):
 
         # KV Connector: update state for finished KV Transfers.
         if model_runner_output.kv_connector_output:
-            self._update_from_kv_xfer_finished(
-                model_runner_output.kv_connector_output)
+            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.
@@ -304,12 +328,14 @@ class DiffusionScheduler(OmniScheduler):
                     eco.finished_requests = finished_set
                 else:
                     engine_core_outputs[client_index] = EngineCoreOutputs(
-                        finished_requests=finished_set)
+                        finished_requests=finished_set
+                    )
             finished_req_ids.clear()
 
         if engine_core_outputs:
             # Return stats to only one of the front-ends.
-            next(iter(engine_core_outputs.values())).scheduler_stats = (
-                self.make_stats(spec_decoding_stats))
+            next(iter(engine_core_outputs.values())).scheduler_stats = self.make_stats(
+                spec_decoding_stats
+            )
 
-        return engine_core_outputs
+        return engine_core_outputs

vllm_omni/core/sched/output.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from vllm.v1.core.sched.output import NewRequestData
-from vllm.v1.request import Request
 from typing import Optional
 
-from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.request import Request
 
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 
 
 @dataclass
@@ -33,4 +33,4 @@ class OmniNewRequestData(NewRequestData):
             lora_request=request.lora_request,
             prompt_embeds=request.prompt_embeds,
             additional_information=request.additional_information,
-        )
+        )

vllm_omni/core/sched/scheduler.py

@@ -1,23 +1,29 @@
 from __future__ import annotations
+
+from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.logger import init_logger
+
 
 class OmniScheduler(VLLMScheduler):
     """
     OmniScheduler: Scheduler for vLLM-omni multimodal processing.
 
-    This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive
-    processing with additional fields and methods specific to vLLM-omni.
+    This scheduler extends vLLM's scheduler to support multimodal and
+    non-autoregressive processing with additional fields and methods
+    specific to vLLM-omni.
     """
 
-    # Ensure scheduled_new_reqs carry omni-specific payloads (e.g., additional_information)
+    # Ensure scheduled_new_reqs carry omni-specific payloads
+    # (e.g., additional_information)
     def schedule(self) -> SchedulerOutput:  # type: ignore[override]
         scheduler_output = super().schedule()
         try:
             # Late import to avoid circulars in some launch modes
             from .output import OmniNewRequestData
-            # Rewrap base NewRequestData entries with OmniNewRequestData, enriching with request-level payloads
+
+            # Rewrap base NewRequestData entries with OmniNewRequestData,
+            # enriching with request-level payloads
             new_list = []
             for nr in scheduler_output.scheduled_new_reqs:
                 req_id = getattr(nr, "req_id", None)
@@ -35,14 +41,22 @@ class OmniScheduler(VLLMScheduler):
                     num_computed_tokens=nr.num_computed_tokens,
                     lora_request=nr.lora_request,
                     # Enrich with omni payloads from the live request object
-                    prompt_embeds=getattr(request, "prompt_embeds", None) if request else None,
-                    additional_information=getattr(request, "additional_information", None) if request else None,
+                    prompt_embeds=(
+                        getattr(request, "prompt_embeds", None) if request else None
+                    ),
+                    additional_information=(
+                        getattr(request, "additional_information", None)
+                        if request
+                        else None
+                    ),
                 )
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
         except Exception:
             # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            init_logger(__name__).exception(
+                "Failed to wrap scheduled_new_reqs with OmniNewRequestData"
+            )
 
-        return scheduler_output
+        return scheduler_output

vllm_omni/engine/__init__.py

@@ -1,11 +1,12 @@
 """
 Engine components for vLLM-omni.
 """
-import msgspec
-from typing import Any
 
+from typing import Any, Optional
+
+import msgspec
 from vllm.v1.engine import EngineCoreRequest
-from typing import Optional
+
 
 class PromptEmbedsPayload(msgspec.Struct):
     """Serialized prompt embeddings payload for direct transfer.
@@ -52,4 +53,3 @@ class OmniEngineCoreRequest(EngineCoreRequest):
     prompt_embeds: Optional[PromptEmbedsPayload] = None
     # Optional additional information dictionary (serialized)
     additional_information: Optional[AdditionalInformationPayload] = None
-    

vllm_omni/engine/arg_utils.py

@@ -1,9 +1,11 @@
-from vllm.engine.arg_utils import EngineArgs
-from typing import Literal, Optional
 from dataclasses import dataclass
-from vllm_omni.config import OmniModelConfig
+from typing import Optional
+
+from vllm.engine.arg_utils import EngineArgs
 from vllm.utils import FlexibleArgumentParser
 
+from vllm_omni.config import OmniModelConfig
+
 
 @dataclass
 class OmniEngineArgs(EngineArgs):
@@ -25,26 +27,34 @@ class OmniEngineArgs(EngineArgs):
                 "model_config.engine_output_type for schedulers to use."
             ),
         )
-        parser.add_argument("--model-stage", type=str, default=OmniEngineArgs.model_stage, 
-        help="Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). This will be written into model_config.model_stage for schedulers to use.")
+        parser.add_argument(
+            "--model-stage",
+            type=str,
+            default=OmniEngineArgs.model_stage,
+            help=(
+                "Declare model stage (e.g., 'thinker', 'talker', 'token2wav'). "
+                "This will be written into model_config.model_stage for "
+                "schedulers to use."
+            ),
+        )
         return parser
 
     def create_model_config(self) -> OmniModelConfig:
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
-        
+
         # Create OmniModelConfig by copying all base config attributes
         # and adding the new omni-specific fields
         config_dict = base_config.__dict__.copy()
-        
+
         # Add the new omni-specific fields
-        config_dict['stage_id'] = self.stage_id
-        config_dict['model_stage'] = self.model_stage
-        config_dict['model_arch'] = self.model_arch
-        config_dict['engine_output_type'] = self.engine_output_type
-        
+        config_dict["stage_id"] = self.stage_id
+        config_dict["model_stage"] = self.model_stage
+        config_dict["model_arch"] = self.model_arch
+        config_dict["engine_output_type"] = self.engine_output_type
+
         # Create and return the OmniModelConfig instance
         omni_config = OmniModelConfig(**config_dict)
         omni_config.hf_config.architectures = omni_config.architectures
-        
-        return omni_config
+
+        return omni_config

vllm_omni/engine/diffusion_engine.py

@@ -5,10 +5,7 @@ from __future__ import annotations
 from typing import Any, Optional
 
 from ..config import DiTConfig
-from ..worker.gpu_diffusion_worker import (
-    DiffusionGPUWorker,
-    DiffusionRunnerOutput,
-)
+from ..worker.gpu_diffusion_worker import DiffusionGPUWorker, DiffusionRunnerOutput
 
 
 class DiffusionEngine:
@@ -27,14 +24,12 @@ class DiffusionEngine:
         config_model_path = None
         if model_cfg is not None:
             if isinstance(model_cfg, dict):
-                config_model_path = (
-                    model_cfg.get("model")
-                    or model_cfg.get("model_path")
+                config_model_path = model_cfg.get("model") or model_cfg.get(
+                    "model_path"
                 )
             else:
-                config_model_path = (
-                    getattr(model_cfg, "model", None)
-                    or getattr(model_cfg, "model_path", None)
+                config_model_path = getattr(model_cfg, "model", None) or getattr(
+                    model_cfg, "model_path", None
                 )
 
         self.model_path = model_path or config_model_path
@@ -116,9 +111,7 @@ class DiffusersPipelineEngine(DiffusionEngine):
             else self.num_inference_steps
         )
         guidance = (
-            float(guidance_scale)
-            if guidance_scale is not None
-            else self.guidance_scale
+            float(guidance_scale) if guidance_scale is not None else self.guidance_scale
         )
 
         return self.worker.generate(
@@ -130,4 +123,4 @@ class DiffusersPipelineEngine(DiffusionEngine):
             negative_prompt=negative_prompt,
             seed=seed,
             image=image,
-        )
+        )

vllm_omni/engine/output_processor.py

@@ -1,21 +1,21 @@
-from typing import Dict, Callable, Optional, Any, Union
+from typing import Any, Callable, Dict, Optional, Union
 
 import torch
-
-from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestState, RequestOutputCollector
-from vllm.transformers_utils.tokenizer_group import TokenizerGroup
-from vllm.v1.engine import FinishReason
-from vllm.v1.metrics.stats import IterationStats
+from vllm.logger import init_logger
 from vllm.sampling_params import RequestOutputKind
 from vllm.transformers_utils.tokenizer import AnyTokenizer
-from vllm.v1.engine import EngineCoreRequest
+from vllm.transformers_utils.tokenizer_group import TokenizerGroup
+from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
+from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.engine.output_processor import (
+    OutputProcessorOutput,
+    RequestOutputCollector,
+    RequestState,
+)
 from vllm.v1.engine.parallel_sampling import ParentRequest
-from vllm.logger import init_logger
-from vllm.v1.engine import EngineCoreOutput
-
+from vllm.v1.metrics.stats import IterationStats
 
 logger = init_logger(__name__)
 
@@ -67,8 +67,9 @@ class OmniRequestState(RequestState):
             request_id=request.request_id,
             parent_req=parent_req,
             request_index=request_index,
-            lora_name=(request.lora_request.name
-                       if request.lora_request is not None else None),
+            lora_name=(
+                request.lora_request.name if request.lora_request is not None else None
+            ),
             output_kind=output_kind,
             prompt=prompt,
             prompt_token_ids=request.prompt_token_ids,
@@ -80,8 +81,9 @@ class OmniRequestState(RequestState):
             log_stats=log_stats,
         )
 
-    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor],
-                               mm_type: Optional[str]) -> None:
+    def add_multimodal_tensor(
+        self, tensor: Optional[torch.Tensor], mm_type: Optional[str]
+    ) -> None:
         if tensor is None:
             return
         try:
@@ -126,36 +128,41 @@ class OmniRequestState(RequestState):
                 pass
 
         request_id = self.request_id
-        output = self._new_completion_output(new_token_ids, finish_reason,
-                                             stop_reason)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
 
         if self.parent_req is None:
             outputs = [output]
         else:
             request_id, outputs, finished = self.parent_req.get_outputs(
-                request_id, output)
+                request_id, output
+            )
             if not outputs:
                 return None
 
-        return self._new_request_output(request_id, outputs, finished,
-                                        kv_transfer_params)
+        return self._new_request_output(
+            request_id, outputs, finished, kv_transfer_params
+        )
 
     def _new_completion_output(
         self,
         token_ids: list[int],
         finish_reason: Optional[FinishReason],
-        stop_reason: Optional[Union[int, str]]
+        stop_reason: Optional[Union[int, str]],
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason,
-                                                     stop_reason)
+        base_output = super()._new_completion_output(
+            token_ids, finish_reason, stop_reason
+        )
         try:
             if self.mm_accumulated is not None:
                 tensor = self.mm_accumulated
                 try:
                     tensor = tensor.detach().to("cpu")
                 except Exception:
-                    logger.debug("Failed to move accumulated multimodal tensor to CPU", exc_info=True)
+                    logger.debug(
+                        "Failed to move accumulated multimodal tensor to CPU",
+                        exc_info=True,
+                    )
                 # Attach on the completion output for downstream consumers.
                 if not hasattr(base_output, "multimodal_output"):
                     setattr(base_output, "multimodal_output", {})
@@ -177,17 +184,24 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
       produce the correct RequestOutput/PoolingRequestOutput.
     - Allow custom per-modality handlers via register_handler().
     """
-    def __init__(self, tokenizer: TokenizerGroup, log_stats: bool, engine_core_output_type: Optional[str] = None):
+
+    def __init__(
+        self,
+        tokenizer: TokenizerGroup,
+        log_stats: bool,
+        engine_core_output_type: Optional[str] = None,
+    ):
         super().__init__(tokenizer=tokenizer, log_stats=log_stats)
         self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {}
         self._reqid_to_mm_type: Dict[str, str] = {}
         self.request_states: dict[str, OmniRequestState] = {}
         self.engine_core_output_type = engine_core_output_type
-        
-    def register_handler(self, modality: str,
-                         handler: Callable[[EngineCoreOutput], None]) -> None:
+
+    def register_handler(
+        self, modality: str, handler: Callable[[EngineCoreOutput], None]
+    ) -> None:
         self.output_handlers[modality.lower()] = handler
-    
+
     def add_request(
         self,
         request: EngineCoreRequest,
@@ -200,21 +214,26 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         if request_id in self.request_states:
             raise ValueError(f"Request id {request_id} already running.")
 
-        tokenizer = None if not self.tokenizer else \
-            self.tokenizer.get_lora_tokenizer(request.lora_request)
+        tokenizer = (
+            None
+            if not self.tokenizer
+            else self.tokenizer.get_lora_tokenizer(request.lora_request)
+        )
 
-        req_state = OmniRequestState.from_new_request(tokenizer=tokenizer,
-                                                  request=request,
-                                                  prompt=prompt,
-                                                  parent_req=parent_req,
-                                                  request_index=request_index,
-                                                  queue=queue,
-                                                  log_stats=self.log_stats)
+        req_state = OmniRequestState.from_new_request(
+            tokenizer=tokenizer,
+            request=request,
+            prompt=prompt,
+            parent_req=parent_req,
+            request_index=request_index,
+            queue=queue,
+            log_stats=self.log_stats,
+        )
         self.request_states[request_id] = req_state
         self.lora_states.add_request(req_state)
         if parent_req:
             self.parent_requests[parent_req.request_id] = parent_req
-    
+
     def process_outputs(
         self,
         engine_core_outputs: list[EngineCoreOutput],
@@ -238,9 +257,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 continue
 
             # 1) Stats
-            self._update_stats_from_output(req_state, eco,
-                                           engine_core_timestamp,
-                                           iteration_stats)
+            self._update_stats_from_output(
+                req_state, eco, engine_core_timestamp, iteration_stats
+            )
 
             new_token_ids = eco.new_token_ids
             pooling_output = eco.pooling_output
@@ -254,7 +273,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             assert req_state.detokenizer is not None
             assert req_state.logprobs_processor is not None
             stop_string = req_state.detokenizer.update(
-                new_token_ids, finish_reason == FinishReason.STOP)
+                new_token_ids, finish_reason == FinishReason.STOP
+            )
             if stop_string:
                 finish_reason = FinishReason.STOP
                 stop_reason = stop_string
@@ -262,11 +282,19 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
             # 2.5) Accumulate multimodal tensors in RequestState
             try:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                if pooling_output is not None and isinstance(req_state, OmniRequestState):
+                mm_type = (
+                    getattr(eco, "output_type", self.engine_core_output_type) or ""
+                ).lower()
+                if pooling_output is not None and isinstance(
+                    req_state, OmniRequestState
+                ):
                     req_state.add_multimodal_tensor(pooling_output, mm_type)
             except Exception:
-                logger.debug("Failed to accumulate multimodal tensor for request %s", req_id, exc_info=True)
+                logger.debug(
+                    "Failed to accumulate multimodal tensor for request %s",
+                    req_id,
+                    exc_info=True,
+                )
 
             # 3) Create RequestOutput objects, forcing combined mode to keep ids
             pooling_for_make = pooling_output
@@ -274,20 +302,29 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 # Do not consume pooling path now; keep ids and attach mm later
                 pooling_for_make = None
 
-            ro = req_state.make_request_output(new_token_ids, pooling_for_make,
-                                               finish_reason, stop_reason,
-                                               kv_transfer_params,
-                                               num_cached_tokens)
+            ro = req_state.make_request_output(
+                new_token_ids,
+                pooling_for_make,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+                num_cached_tokens,
+            )
             if ro:
                 # Attach accumulated multimodal payload if any
                 try:
-                    if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
+                    if (
+                        isinstance(req_state, OmniRequestState)
+                        and req_state.mm_accumulated is not None
+                    ):
                         mm_key = req_state.mm_type or "latents"
                         if not hasattr(ro, "multimodal_output"):
                             setattr(ro, "multimodal_output", {})
                         ro.multimodal_output[mm_key] = req_state.mm_accumulated
                 except Exception:
-                    logger.exception("Error attaching multimodal payload in process_outputs")
+                    logger.exception(
+                        "Error attaching multimodal payload in process_outputs"
+                    )
                 if req_state.queue is not None:
                     req_state.queue.put(ro)
                 else:
@@ -301,8 +338,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self.parent_requests.pop(parent_req.request_id, None)
                 if not eco.finished:
                     reqs_to_abort.append(req_id)
-                self._update_stats_from_finished(req_state, finish_reason,
-                                                 iteration_stats)
+                self._update_stats_from_finished(
+                    req_state, finish_reason, iteration_stats
+                )
                 # Cleanup per-request mm state
                 if isinstance(req_state, OmniRequestState):
                     req_state.mm_accumulated = None
@@ -312,10 +350,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             request_outputs=request_outputs,
             reqs_to_abort=reqs_to_abort,
         )
-    
+
     # ---- routing helpers ----
     def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
-        output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
+        output_type = (
+            getattr(eco, "output_type", self.engine_core_output_type) or ""
+        ).lower()
 
         # Custom handler first (if registered)
         if output_type in self.output_handlers:
@@ -347,7 +387,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure image tensors are surfaced via pooling_output for vLLM."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("image", "images", "pixel_values", "pixels"))
+                eco, keys=("image", "images", "pixel_values", "pixels")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -357,8 +398,17 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         # Preserve text tokens as-is; ensure pooling_output carries image/latents
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("image", "images", "pixel_values", "pixels",
-                           "latent", "latents", "z"))
+                eco,
+                keys=(
+                    "image",
+                    "images",
+                    "pixel_values",
+                    "pixels",
+                    "latent",
+                    "latents",
+                    "z",
+                ),
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -366,7 +416,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure latent tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("latent", "latents", "z", "posterior"))
+                eco, keys=("latent", "latents", "z", "posterior")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -374,8 +425,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         """Ensure audio tensors are surfaced via pooling_output."""
         if eco.pooling_output is None:
             tensor = self._extract_from_multimodal_outputs(
-                eco, keys=("audio", "audios", "wav", "waveform",
-                           "audio_pcm", "pcm"))
+                eco, keys=("audio", "audios", "wav", "waveform", "audio_pcm", "pcm")
+            )
             if tensor is not None:
                 eco.pooling_output = tensor
 
@@ -408,4 +459,4 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         for v in mm.values():
             if isinstance(v, torch.Tensor):
                 return v
-        return None
+        return None

vllm_omni/engine/processor.py

@@ -2,23 +2,28 @@ import time
 from collections.abc import Mapping, Sequence
 from typing import Any, Optional, Union
 
+import torch
+from vllm.config import VllmConfig
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.inputs.parse import split_enc_dec_inputs
 from vllm.lora.request import LoRARequest
-from vllm.multimodal import MultiModalKwargs
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs, MultiModalRegistry
 from vllm.multimodal.inputs import PlaceholderRange
 from vllm.multimodal.utils import merge_and_sort_multimodal_metadata
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
-from vllm.config import VllmConfig
 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-from vllm.platforms import current_platform
-
 from vllm.v1.engine.processor import Processor
-from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload, AdditionalInformationEntry, OmniEngineCoreRequest
-import torch
+
+from vllm_omni.engine import (
+    AdditionalInformationEntry,
+    AdditionalInformationPayload,
+    OmniEngineCoreRequest,
+    PromptEmbedsPayload,
+)
+from vllm_omni.inputs.preprocess import OmniInputPreprocessor
+
 
 class OmniProcessor(Processor):
     @staticmethod
@@ -42,15 +47,17 @@ class OmniProcessor(Processor):
             torch.bool: "bool",
         }
         return mapping.get(dtype, str(dtype).replace("torch.", ""))
-    def __init__(self, 
-                 vllm_config: VllmConfig,
-                 tokenizer: TokenizerGroup,
-                 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-                 ):
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        tokenizer: TokenizerGroup,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+    ):
         super().__init__(vllm_config, tokenizer, mm_registry)
-        self.input_preprocessor = OmniInputPreprocessor(self.model_config,
-                                                    self.tokenizer,
-                                                    mm_registry)
+        self.input_preprocessor = OmniInputPreprocessor(
+            self.model_config, self.tokenizer, mm_registry
+        )
 
     def process_inputs(
         self,
@@ -73,10 +80,13 @@ class OmniProcessor(Processor):
             raise ValueError("V1 does not support tracing yet.")
 
         data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank <
-                                                   data_parallel_size):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} "
-                             f"is out of range [0, {data_parallel_size}).")
+        if data_parallel_rank is not None and not (
+            0 <= data_parallel_rank < data_parallel_size
+        ):
+            raise ValueError(
+                f"data_parallel_rank {data_parallel_rank} "
+                f"is out of range [0, {data_parallel_size})."
+            )
 
         if arrival_time is None:
             arrival_time = time.time()
@@ -85,8 +95,9 @@ class OmniProcessor(Processor):
         # 1. Tokenize text prompt, with LoRA request if one exists.
         # 2. For multimodal models with a merged preprocessor, preprocess
         #   multimodal data and expand prompt token ids accordingly.
-        return_mm_hashes = (self.model_config.processor_return_mm_hashes
-                            or bool(self.cache_config.enable_prefix_caching))
+        return_mm_hashes = self.model_config.processor_return_mm_hashes or bool(
+            self.cache_config.enable_prefix_caching
+        )
         processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
             prompt,
             tokenization_kwargs=tokenization_kwargs,
@@ -115,14 +126,16 @@ class OmniProcessor(Processor):
             sampling_params = params.clone()
             # If unset max tokens, then generate up to the max_model_len.
             if sampling_params.max_tokens is None:
-                sampling_params.max_tokens = (
-                    self.model_config.max_model_len -
-                    len(decoder_inputs["prompt_token_ids"]))
+                sampling_params.max_tokens = self.model_config.max_model_len - len(
+                    decoder_inputs["prompt_token_ids"]
+                )
             sampling_params.update_from_generation_config(
-                self.generation_config_fields, eos_token_id)
+                self.generation_config_fields, eos_token_id
+            )
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(
-                    self.tokenizer.get_lora_tokenizer(lora_request))
+                    self.tokenizer.get_lora_tokenizer(lora_request)
+                )
         else:
             pooling_params = params.clone()
 
@@ -159,22 +172,23 @@ class OmniProcessor(Processor):
                     items = decoder_mm_inputs.get_items(modality)
                     item = items[used_indices[modality]]
 
-                    orig_sorted_mm_inputs.append(
-                        MultiModalKwargs.from_items([item]))
+                    orig_sorted_mm_inputs.append(MultiModalKwargs.from_items([item]))
                     used_indices[modality] += 1
             else:
                 orig_sorted_mm_inputs = [
-                    MultiModalKwargs.from_items([item]) for item in
-                    decoder_mm_inputs.get_items(sorted_item_modalities[0])
+                    MultiModalKwargs.from_items([item])
+                    for item in decoder_mm_inputs.get_items(sorted_item_modalities[0])
                 ]
 
             if sorted_mm_hashes is not None:
                 sorted_mm_inputs = self.mm_input_cache_client.get_and_update(
-                    orig_sorted_mm_inputs, sorted_mm_hashes)
+                    orig_sorted_mm_inputs, sorted_mm_hashes
+                )
             else:
                 sorted_mm_inputs = orig_sorted_mm_inputs
-        
-        # Serialize prompt_embeds and additional_information if provided (direct-transfer path)
+
+        # Serialize prompt_embeds and additional_information if provided
+        # (direct-transfer path)
         prompt_embeds_payload: Optional[PromptEmbedsPayload] = None
         additional_information_payload: Optional[AdditionalInformationPayload] = None
         if "prompt_embeds" in decoder_inputs:  # type: ignore[operator]
@@ -182,7 +196,8 @@ class OmniProcessor(Processor):
             pe: torch.Tensor = decoder_inputs["prompt_embeds"]  # type: ignore[index]
             if pe.ndim != 2:
                 raise ValueError(
-                    "prompt_embeds must be of shape (seq_len, hidden_size)")
+                    "prompt_embeds must be of shape (seq_len, hidden_size)"
+                )
             # Move to CPU and ensure contiguous memory for stable serialization
             pe_cpu = pe.detach().to("cpu").contiguous()
             seq_len, hidden_size = pe_cpu.shape
@@ -194,7 +209,7 @@ class OmniProcessor(Processor):
                 dtype=dtype_str,
             )
         if "additional_information" in decoder_inputs:  # type: ignore[operator]
-            raw_info: dict[str, Any] = decoder_inputs["additional_information"]  # type: ignore[index]
+            raw_info: dict[str, Any] = decoder_inputs["additional_information"]  # type: ignore[index]  # noqa: E501
             entries: dict[str, AdditionalInformationEntry] = {}
             for key, value in raw_info.items():
                 if isinstance(value, torch.Tensor):
@@ -210,10 +225,12 @@ class OmniProcessor(Processor):
                     entry = AdditionalInformationEntry(list_data=value)
                 else:
                     raise ValueError(
-                        "additional_information values must be Tensor or list")
+                        "additional_information values must be Tensor or list"
+                    )
                 entries[key] = entry
             additional_information_payload = AdditionalInformationPayload(
-                entries=entries)
+                entries=entries
+            )
 
         return decoder_inputs.get("prompt"), OmniEngineCoreRequest(
             request_id=request_id,
@@ -231,4 +248,4 @@ class OmniProcessor(Processor):
             data_parallel_rank=data_parallel_rank,
             prompt_embeds=prompt_embeds_payload,
             additional_information=additional_information_payload,
-        )
+        )

vllm_omni/entrypoints/api_server.py

@@ -3,17 +3,18 @@ API server for vLLM-omni.
 """
 
 import asyncio
-from typing import Dict, Any, List
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
+from typing import Any, Dict, List
+
 import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 
 from .omni_llm import AsyncOmniLLM
 
 
 class GenerateRequest(BaseModel):
     """Request model for generation."""
+
     prompts: List[str]
     max_tokens: int = 100
     temperature: float = 0.7
@@ -22,6 +23,7 @@ class GenerateRequest(BaseModel):
 
 class GenerateResponse(BaseModel):
     """Response model for generation."""
+
     outputs: List[Dict[str, Any]]
     stage_outputs: List[Dict[str, Any]] = None
 
@@ -29,7 +31,7 @@ class GenerateResponse(BaseModel):
 app = FastAPI(
     title="vLLM-omni API",
     description="Multi-modality models inference and serving",
-    version="0.1.0"
+    version="0.1.0",
 )
 
 # Global omni_llm instance
@@ -50,48 +52,51 @@ async def generate(request: GenerateRequest):
     try:
         if omni_llm is None:
             raise HTTPException(status_code=500, detail="OmniLLM not initialized")
-        
+
         # Prepare stage arguments
         if request.stage_args is None:
             # Create default stage arguments - one per stage config
             # For now, we'll process all prompts in the first stage
-            stage_args = [{
-                "prompt": " ".join(request.prompts) if request.prompts else "",
-                "max_tokens": request.max_tokens,
-                "temperature": request.temperature
-            }]
+            stage_args = [
+                {
+                    "prompt": " ".join(request.prompts) if request.prompts else "",
+                    "max_tokens": request.max_tokens,
+                    "temperature": request.temperature,
+                }
+            ]
         else:
             stage_args = request.stage_args
-        
+
         # Generate using omni_llm
         try:
             outputs = await omni_llm.generate_async(stage_args)
-        except Exception as e:
+        except Exception:
             # Fallback to synchronous generation
             outputs = omni_llm.generate(stage_args)
-        
+
         # Convert outputs to response format
         response_outputs = []
         for output in outputs:
-            if hasattr(output, 'outputs') and output.outputs:
+            if hasattr(output, "outputs") and output.outputs:
                 for out in output.outputs:
-                    response_outputs.append({
-                        "text": getattr(out, 'text', ''),
-                        "finished": getattr(out, 'finish_reason', 'length') != 'length',
-                        "tokens": getattr(out, 'token_ids', [])
-                    })
+                    response_outputs.append(
+                        {
+                            "text": getattr(out, "text", ""),
+                            "finished": getattr(out, "finish_reason", "length")
+                            != "length",
+                            "tokens": getattr(out, "token_ids", []),
+                        }
+                    )
             else:
-                response_outputs.append({
-                    "text": "",
-                    "finished": True,
-                    "tokens": []
-                })
-        
+                response_outputs.append({"text": "", "finished": True, "tokens": []})
+
         return GenerateResponse(
             outputs=response_outputs,
-            stage_outputs=[{"stage": i, "output": "processed"} for i in range(len(stage_args))]
+            stage_outputs=[
+                {"stage": i, "output": "processed"} for i in range(len(stage_args))
+            ],
         )
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -107,36 +112,35 @@ async def get_info():
     """Get information about the service."""
     if omni_llm is None:
         return {"error": "OmniLLM not initialized"}
-    
+
     return {
         "service": "vllm-omni",
         "version": "0.1.0",
-        "num_stages": omni_llm.get_num_stages() if hasattr(omni_llm, 'get_num_stages') else 0,
+        "num_stages": (
+            omni_llm.get_num_stages() if hasattr(omni_llm, "get_num_stages") else 0
+        ),
         "stage_configs": [
             {
                 "stage_id": config.stage_id,
                 "engine_type": config.engine_type,
                 "model_path": config.model_path,
                 "input_modalities": config.input_modalities,
-                "output_modalities": config.output_modalities
+                "output_modalities": config.output_modalities,
             }
             for config in omni_llm.stage_configs
-        ]
+        ],
     }
 
 
-async def run_server(omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000):
+async def run_server(
+    omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000
+):
     """Run the API server."""
     global omni_llm
     omni_llm = omni_llm_instance
-    
-    config = uvicorn.Config(
-        app=app,
-        host=host,
-        port=port,
-        log_level="info"
-    )
-    
+
+    config = uvicorn.Config(app=app, host=host, port=port, log_level="info")
+
     server = uvicorn.Server(config)
     await server.serve()
 

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(
-            "--ar-stage",
-            type=str,
-            help="AR stage model path"
-        )
-        
-        parser.add_argument(
-            "--dit-stage",
-            type=str,
-            help="DiT stage model path"
-        )
-        
+        parser.add_argument("--ar-stage", type=str, help="AR stage model path")
+
+        parser.add_argument("--dit-stage", type=str, help="DiT stage model path")
+
         parser.add_argument(
             "--dit-steps",
             type=int,
             default=50,
-            help="Number of DiT inference steps"
+            help="Number of DiT inference steps",
         )
-        
+
         parser.add_argument(
             "--dit-guidance-scale",
             type=float,
             default=7.5,
-            help="DiT guidance scale"
+            help="DiT guidance scale",
         )
-        
+
         parser.add_argument(
             "--use-diffusers",
             action="store_true",
-            help="Use diffusers pipeline for DiT stage"
+            help="Use diffusers pipeline for DiT stage",
         )
-        
+
         # Other arguments
         parser.add_argument(
-            "--log-stats",
-            action="store_true",
-            help="Enable logging statistics"
+            "--log-stats", action="store_true", help="Enable logging statistics"
         )
-        
+
         return parser
-    
+
     def run(self, args: List[str]) -> None:
         """Run the omni serve command."""
         parsed_args = self.parser.parse_args(args)
-        
+
         # Create stage configurations
         stage_configs = self._create_stage_configs(parsed_args)
-        
+
         # Create AsyncOmniLLM instance
         omni_llm = AsyncOmniLLM(
-            stage_configs=stage_configs,
-            log_stats=parsed_args.log_stats
+            stage_configs=stage_configs, log_stats=parsed_args.log_stats
         )
-        
+
         # Start the server
         asyncio.run(self._start_server(omni_llm, parsed_args))
-    
+
     def _create_stage_configs(self, args) -> List:
         """Create stage configurations based on arguments."""
         stage_configs = []
         stage_id = 0
-        
+
         # Add AR stage - use main model if ar_stage not specified
         ar_model = args.ar_stage if args.ar_stage else args.model
         ar_config = create_ar_stage_config(
             stage_id=stage_id,
             model_path=ar_model,
             input_modalities=["text"],
-            output_modalities=["text"]
+            output_modalities=["text"],
         )
         stage_configs.append(ar_config)
         stage_id += 1
-        
+
         # Add DiT stage if specified
         if args.dit_stage:
             dit_model = args.dit_stage
@@ -131,7 +117,7 @@ class OmniServeCommand:
             dit_model = "stabilityai/stable-diffusion-2-1"
         else:
             dit_model = None
-            
+
         if dit_model:
             dit_config = DiTConfig(
                 num_inference_steps=args.dit_steps,
@@ -139,17 +125,17 @@ class OmniServeCommand:
                 use_diffusers=args.use_diffusers,
                 diffusers_pipeline="auto" if args.use_diffusers else None,
             )
-            
+
             dit_stage_config = create_dit_stage_config(
                 stage_id=stage_id,
                 model_path=dit_model,
                 input_modalities=["text"],
                 output_modalities=["image"],
-                dit_config=dit_config
+                dit_config=dit_config,
             )
             stage_configs.append(dit_stage_config)
             stage_id += 1
-        
+
         # If no specific stages are specified, use the main model
         if not stage_configs:
             # Try to detect if it's a multimodal model
@@ -159,23 +145,23 @@ class OmniServeCommand:
                     stage_id=0,
                     model_path=args.model,
                     input_modalities=["text"],
-                    output_modalities=["text"]
+                    output_modalities=["text"],
                 )
                 stage_configs.append(ar_config)
-                
+
                 dit_config = DiTConfig(
                     num_inference_steps=args.dit_steps,
                     guidance_scale=args.dit_guidance_scale,
                     use_diffusers=args.use_diffusers,
                     diffusers_pipeline="auto" if args.use_diffusers else None,
                 )
-                
+
                 dit_stage_config = create_dit_stage_config(
                     stage_id=1,
                     model_path=args.model,
                     input_modalities=["text"],
                     output_modalities=["image"],
-                    dit_config=dit_config
+                    dit_config=dit_config,
                 )
                 stage_configs.append(dit_stage_config)
             else:
@@ -184,23 +170,19 @@ class OmniServeCommand:
                     stage_id=0,
                     model_path=args.model,
                     input_modalities=["text"],
-                    output_modalities=["text"]
+                    output_modalities=["text"],
                 )
                 stage_configs.append(ar_config)
-        
+
         return stage_configs
-    
+
     async def _start_server(self, omni_llm: AsyncOmniLLM, args) -> None:
         """Start the API server."""
         try:
             # Import here to avoid circular imports
             from vllm_omni.entrypoints.api_server import run_server
-            
-            await run_server(
-                omni_llm_instance=omni_llm,
-                host=args.host,
-                port=args.port
-            )
+
+            await run_server(omni_llm_instance=omni_llm, host=args.host, port=args.port)
         except KeyboardInterrupt:
             print("\nShutting down server...")
         except Exception as e:

vllm_omni/entrypoints/omni_llm.py

@@ -1,88 +1,105 @@
-from typing import Union, Sequence, Optional, Any
+from typing import Any, Optional, Sequence, Union
+
 import cloudpickle
-from pydantic import ValidationError
-
-from vllm.inputs import PromptType
-from vllm.sampling_params import SamplingParams
-from vllm.entrypoints.llm import LLM
-
-from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.engine.arg_utils import HfOverrides
-from vllm.usage.usage_lib import UsageContext
-from vllm.config import CompilationConfig, is_init_field
-from vllm.utils import Counter
-from vllm.logger import init_logger
 import vllm.envs as envs
+from pydantic import ValidationError
+from vllm.config import CompilationConfig, is_init_field
+from vllm.engine.arg_utils import HfOverrides
+from vllm.entrypoints.llm import LLM
+from vllm.inputs import PromptType
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import Counter
+from vllm.v1.engine.llm_engine import LLMEngine
 
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model
-from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
+from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.utils import load_stage_configs_from_model
 from vllm_omni.outputs import OmniRequestOutput
 
-
 logger = init_logger(__name__)
 
 
 class OmniLLM:
-    def __init__(self, model: str, stage_configs = None, log_stats: bool = False, **kwargs):
+    def __init__(
+        self, model: str, stage_configs=None, log_stats: bool = False, **kwargs
+    ):
         if stage_configs is None:
             self.initialize_stage_configs(model)
         else:
             self.stage_configs = stage_configs
-        
+
         self.stage_list = []
         self.initialize_stages(model)
-        
+
     def initialize_stage_configs(self, model: str):
         self.stage_configs = load_stage_configs_from_model(model)
-    
+
     def initialize_stages(self, model: str):
         for stage_config in self.stage_configs:
             stage = OmniStage(stage_config)
             stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
             stage.set_engine(stage_llm)
             self.stage_list.append(stage)
-    
+
     def generate(
         self,
         prompts: Union[PromptType, Sequence[PromptType]],
-        sampling_params_list: Optional[Union[SamplingParams,
-                                        Sequence[SamplingParams]]] = None,
+        sampling_params_list: Optional[
+            Union[SamplingParams, Sequence[SamplingParams]]
+        ] = None,
     ) -> list[OmniRequestOutput]:
         """Generate text outputs for the given prompts."""
         final_outputs: list[OmniRequestOutput] = []
         if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
+            raise ValueError(
+                f"Expected {len(self.stage_list)} sampling params, "
+                f"got {len(sampling_params_list)}"
+            )
         for stage_id, stage in enumerate(self.stage_list):
             if stage_id > 0:
                 engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
             else:
                 engine_inputs = prompts
-            engine_outputs = self._run_generation(stage, sampling_params_list[stage_id], engine_inputs)
+            engine_outputs = self._run_generation(
+                stage, sampling_params_list[stage_id], engine_inputs
+            )
             stage.set_engine_outputs(engine_outputs)
-            if hasattr(stage, 'final_output') and stage.final_output:
-                final_outputs.append(OmniRequestOutput(
-                    stage_id=stage_id, 
-                    final_output_type=stage.final_output_type, 
-                    request_output=engine_outputs))
+            if hasattr(stage, "final_output") and stage.final_output:
+                final_outputs.append(
+                    OmniRequestOutput(
+                        stage_id=stage_id,
+                        final_output_type=stage.final_output_type,
+                        request_output=engine_outputs,
+                    )
+                )
         return final_outputs
-    
-    def _run_generation(self, stage: OmniStage, sampling_params: SamplingParams, prompts: Union[PromptType, Sequence[PromptType]]):
+
+    def _run_generation(
+        self,
+        stage: OmniStage,
+        sampling_params: SamplingParams,
+        prompts: Union[PromptType, Sequence[PromptType]],
+    ):
         engine_outputs = []
         for ro in stage.engine.generate(prompts, sampling_params):
             engine_outputs.append(ro)
         return engine_outputs
 
 
-class OmniStageLLM(LLM): 
-    def __init__(self, 
-                model: str, 
-                compilation_config: Optional[Union[int, dict[str, Any],
-                                           CompilationConfig]] = None,
-                hf_overrides: Optional[HfOverrides] = None,
-                **kwargs):
+class OmniStageLLM(LLM):
+    def __init__(
+        self,
+        model: str,
+        compilation_config: Optional[
+            Union[int, dict[str, Any], CompilationConfig]
+        ] = None,
+        hf_overrides: Optional[HfOverrides] = None,
+        **kwargs,
+    ):
         """LLM constructor."""
         if "disable_log_stats" not in kwargs:
             kwargs["disable_log_stats"] = True
@@ -95,21 +112,23 @@ class OmniStageLLM(LLM):
                 kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
 
         if "kv_transfer_config" in kwargs and isinstance(
-                kwargs["kv_transfer_config"], dict):
+            kwargs["kv_transfer_config"], dict
+        ):
             from vllm.config import KVTransferConfig
+
             raw_config_dict = kwargs["kv_transfer_config"]
             try:
-                kwargs["kv_transfer_config"] = KVTransferConfig(
-                    **raw_config_dict)
+                kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
             except ValidationError as e:
                 logger.error(
                     "Failed to convert 'kv_transfer_config' dict to "
                     "KVTransferConfig object. Dict: %s. Error: %s",
-                    raw_config_dict, e)
+                    raw_config_dict,
+                    e,
+                )
                 # Consider re-raising a more specific vLLM error or ValueError
                 # to provide better context to the user.
-                raise ValueError(
-                    f"Invalid 'kv_transfer_config' provided: {e}") from e
+                raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
 
         if hf_overrides is None:
             hf_overrides = {}
@@ -117,11 +136,16 @@ class OmniStageLLM(LLM):
         if compilation_config is not None:
             if isinstance(compilation_config, int):
                 compilation_config_instance = CompilationConfig(
-                    level=compilation_config)
+                    level=compilation_config
+                )
             elif isinstance(compilation_config, dict):
-                predicate = lambda x: is_init_field(CompilationConfig, x[0])
+
+                def predicate(x):
+                    return is_init_field(CompilationConfig, x[0])
+
                 compilation_config_instance = CompilationConfig(
-                    **dict(filter(predicate, compilation_config.items())))
+                    **dict(filter(predicate, compilation_config.items()))
+                )
             else:
                 compilation_config_instance = compilation_config
         else:
@@ -136,22 +160,24 @@ class OmniStageLLM(LLM):
 
         # Create the Engine (autoselects V0 vs V1)
         self.llm_engine = LLMEngine.from_engine_args(
-            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
-        self.llm_engine.output_processor = MultimodalOutputProcessor(tokenizer=self.llm_engine.tokenizer, 
-                                                                    log_stats=self.llm_engine.log_stats)
-        self.llm_engine.processor = OmniProcessor(vllm_config=self.llm_engine.vllm_config,
-                                                  tokenizer=self.llm_engine.tokenizer)
+            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
+        )
+        self.llm_engine.output_processor = MultimodalOutputProcessor(
+            tokenizer=self.llm_engine.tokenizer, log_stats=self.llm_engine.log_stats
+        )
+        self.llm_engine.processor = OmniProcessor(
+            vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
+        )
         self.engine_class = type(self.llm_engine)
 
         self.request_counter = Counter()
         self.default_sampling_params: Union[dict[str, Any], None] = None
 
         if envs.VLLM_USE_V1:
-            supported_tasks = self.llm_engine \
-                .get_supported_tasks()  # type: ignore
+            supported_tasks = self.llm_engine.get_supported_tasks()  # type: ignore
         else:
             supported_tasks = self.llm_engine.model_config.supported_tasks
 
         logger.info("Supported_tasks: %s", supported_tasks)
 
-        self.supported_tasks = supported_tasks
+        self.supported_tasks = supported_tasks

vllm_omni/entrypoints/omni_stage.py

@@ -4,10 +4,11 @@ Stage manager for orchestrating multiple engines in vLLM-omni.
 
 import importlib
 from typing import List, Union
-from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.v1.engine.async_llm import AsyncLLM
+
 from vllm.inputs import TextPrompt
 from vllm.v1.engine import EngineCoreOutput
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.llm_engine import LLMEngine
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
@@ -20,26 +21,28 @@ class OmniStage:
         self.stage_id = stage_config.stage_id
         self.engine_args = stage_config.engine_args
         self.model_stage = stage_config.engine_args.model_stage
-        if hasattr(stage_config, 'engine_input_source'):
+        if hasattr(stage_config, "engine_input_source"):
             self.engine_input_source = stage_config.engine_input_source
         else:
             self.engine_input_source = []
         self.engine_output_type = stage_config.engine_args.engine_output_type
         self.engine_outputs = None
-        if hasattr(stage_config, 'custom_process_input_func'):
+        if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
-            module_path, func_name = stage_config.custom_process_input_func.rsplit('.', 1)
+            module_path, func_name = stage_config.custom_process_input_func.rsplit(
+                ".", 1
+            )
             module = importlib.import_module(module_path)
             self.custom_process_input_func = getattr(module, func_name)
         else:
             self.custom_process_input_func = None
 
-        if hasattr(stage_config, 'final_output'):
+        if hasattr(stage_config, "final_output"):
             self.final_output = stage_config.final_output
         else:
             self.final_output = False
 
-        if hasattr(stage_config, 'final_output_type'):
+        if hasattr(stage_config, "final_output_type"):
             self.final_output_type = stage_config.final_output_type
         else:
             self.final_output_type = None
@@ -51,12 +54,14 @@ class OmniStage:
     def set_async_engine(self, async_engine: AsyncLLM) -> None:
         """Initialize the async engine for the stage."""
         self.async_engine = async_engine
-    
+
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
         """Set the engine output for the stage."""
         self.engine_outputs = engine_outputs
-    
-    def process_engine_inputs(self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None) -> List[Union[OmniTokensPrompt, TextPrompt]]:
+
+    def process_engine_inputs(
+        self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+    ) -> List[Union[OmniTokensPrompt, TextPrompt]]:
         """Process the engine input for the stage."""
         if self.custom_process_input_func is None:
             engine_inputs = []
@@ -64,17 +69,25 @@ class OmniStage:
                 raise ValueError("engine_input_source is empty")
             source_stage_id = self.engine_input_source[0]
             source_outputs = stage_list[source_stage_id].engine_outputs
-            multi_modal_data = {source_output.request_id: 
-            prompt.get('multi_modal_data', None) for source_output, prompt in zip(source_outputs, prompt)}
+            multi_modal_data = {
+                source_output.request_id: prompt.get("multi_modal_data", None)
+                for source_output, prompt in zip(source_outputs, prompt)
+            }
 
             for source_output in source_outputs:
                 engine_input = OmniTokensPrompt(
-                    prompt_token_ids = source_output.outputs[0].token_ids,
-                    multi_modal_data=multi_modal_data[source_output.request_id] if multi_modal_data else None,
+                    prompt_token_ids=source_output.outputs[0].token_ids,
+                    multi_modal_data=(
+                        multi_modal_data[source_output.request_id]
+                        if multi_modal_data
+                        else None
+                    ),
                 )
                 engine_inputs.append(engine_input)
             return engine_inputs
 
         else:
             engine_input_source = self.engine_input_source
-            return self.custom_process_input_func(stage_list, engine_input_source, prompt)
+            return self.custom_process_input_func(
+                stage_list, engine_input_source, prompt
+            )

vllm_omni/entrypoints/utils.py

@@ -1,5 +1,6 @@
 import os
 from pathlib import Path
+
 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config
 
@@ -22,4 +23,4 @@ def load_stage_configs_from_model(model: str):
 def load_stage_configs_from_yaml(config_path: str):
     """Load stage configs from yaml file."""
     config_data = OmegaConf.load(config_path)
-    return config_data.stage_args
+    return config_data.stage_args

vllm_omni/inputs/data.py

@@ -1,6 +1,7 @@
-from vllm.inputs.data import TokensPrompt, EmbedsPrompt, TokenInputs
 from typing import Any, NotRequired, Optional
+
 import torch
+from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
 
 
 class OmniTokensPrompt(TokensPrompt):
@@ -53,4 +54,4 @@ def token_inputs_omni(
     if additional_information is not None:
         inputs["additional_information"] = additional_information
 
-    return inputs
+    return inputs

vllm_omni/inputs/parse.py

@@ -1,5 +1,11 @@
-from vllm.inputs.parse import ParsedSingletonPrompt, ParsedStrPrompt, ParsedEmbedsPrompt, ParsedTokensPrompt, ParsedTextPrompt
 from vllm.inputs.data import SingletonPrompt
+from vllm.inputs.parse import (
+    ParsedEmbedsPrompt,
+    ParsedSingletonPrompt,
+    ParsedStrPrompt,
+    ParsedTextPrompt,
+    ParsedTokensPrompt,
+)
 
 
 def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPrompt:
@@ -8,14 +14,18 @@ def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPromp
     elif isinstance(prompt, dict):
         # Type ignores are because mypy does not correctly infer the TypedDicts
         # Pyright does succeed.
-        # Priority tokens: When both tokens and embeds exist, keep both and follow the tokens path
+        # Priority tokens: When both tokens and embeds exist, keep both and
+        # follow the tokens path
         if "prompt_token_ids" in prompt:
             return ParsedTokensPrompt(
-                type="tokens", content=prompt)  # type: ignore[typeddict-item]
+                type="tokens", content=prompt
+            )  # type: ignore[typeddict-item]
         elif "prompt_embeds" in prompt:
             return ParsedEmbedsPrompt(
-                type="embeds", content=prompt)  # type: ignore[typeddict-item]
+                type="embeds", content=prompt
+            )  # type: ignore[typeddict-item]
         elif "prompt" in prompt:
             return ParsedTextPrompt(type="text", content=prompt)
     raise TypeError(
-        "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")
+        "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt"
+    )

vllm_omni/inputs/preprocess.py

@@ -1,10 +1,11 @@
-from typing import Optional, Any, Union
-from typing_extensions import assert_never
+from typing import Any, Optional, Union
 
-from vllm.lora.request import LoRARequest
+from typing_extensions import assert_never
+from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.preprocess import InputPreprocessor
-from vllm.inputs.data import TokensPrompt, SingletonPrompt, SingletonInputs, TextPrompt
+from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import MultiModalInputs
+
 from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
 from vllm_omni.inputs.parse import parse_singleton_prompt_omni
 
@@ -44,7 +45,7 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs["cache_salt"] = cache_salt
 
         return inputs
-    
+
     def _prompt_to_llm_inputs(
         self,
         prompt: SingletonPrompt,
@@ -91,7 +92,7 @@ class OmniInputPreprocessor(InputPreprocessor):
             )
 
         assert_never(parsed)
-    
+
     async def _process_tokens_async(
         self,
         parsed_content: TokensPrompt,
@@ -126,7 +127,7 @@ class OmniInputPreprocessor(InputPreprocessor):
             inputs["cache_salt"] = cache_salt
 
         return inputs
-    
+
     async def _prompt_to_llm_inputs_async(
         self,
         prompt: SingletonPrompt,
@@ -163,4 +164,4 @@ class OmniInputPreprocessor(InputPreprocessor):
                 return_mm_hashes=return_mm_hashes,
             )
 
-        assert_never(parsed)
+        assert_never(parsed)

vllm_omni/model_executor/layers/mrope.py

@@ -4,9 +4,8 @@ from typing import Optional, Union
 import numpy as np
 import torch
 from transformers import PretrainedConfig
-
 from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
-from vllm.model_executor.layers.rotary_embedding.common import apply_rotary_emb_dispatch
+
 
 def _apply_rotary_emb(
     x: torch.Tensor,
@@ -36,6 +35,7 @@ def _apply_rotary_emb(
     else:
         return torch.stack((o1, o2), dim=-1).flatten(-2)
 
+
 class MRotaryEmbedding(RotaryEmbedding):
     """Rotary Embedding with Multimodal Sections."""
 
@@ -53,12 +53,21 @@ class MRotaryEmbedding(RotaryEmbedding):
         # the input video. We enlarge max_position_embeddings to 4 times to get
         # a larger the cos and sin cache.
         self.cache_max_position_num = max_position_embeddings * 4
-        super().__init__(head_size, rotary_dim, self.cache_max_position_num,
-                         base, is_neox_style, dtype)
+        super().__init__(
+            head_size,
+            rotary_dim,
+            self.cache_max_position_num,
+            base,
+            is_neox_style,
+            dtype,
+        )
 
         self.mrope_section = mrope_section
         if self.mrope_section:
-            print("Warning: mrope_section check is disabled in Qwen2.5-Omni, this may cause errors, and should be restored in the future.")
+            print(
+                "Warning: mrope_section check is disabled in Qwen2.5-Omni, "
+                "this may cause errors, and should be restored in the future."
+            )
             # assert sum(self.mrope_section) == rotary_dim // 2
 
     def forward(
@@ -85,35 +94,31 @@ class MRotaryEmbedding(RotaryEmbedding):
         if positions.ndim == 2:
             assert self.mrope_section
 
-            cos = torch.cat([
-                m[i]
-                for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
-            ],
-                            dim=-1)
-            sin = torch.cat([
-                m[i]
-                for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
-            ],
-                            dim=-1)
+            cos = torch.cat(
+                [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
+            sin = torch.cat(
+                [m[i] for i, m in enumerate(sin.split(self.mrope_section, dim=-1))],
+                dim=-1,
+            )
 
         query_shape = query.shape
         query = query.view(num_tokens, -1, self.head_size)
-        query_rot = query[..., :self.rotary_dim]
-        query_pass = query[..., self.rotary_dim:]
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
         # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
         #                                       self.is_neox_style)
-        query_rot = _apply_rotary_emb(query_rot, cos, sin,
-                                              self.is_neox_style)
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
         query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
 
         key_shape = key.shape
         key = key.view(num_tokens, -1, self.head_size)
-        key_rot = key[..., :self.rotary_dim]
-        key_pass = key[..., self.rotary_dim:]
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
         # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
         #                                     self.is_neox_style)
-        key_rot = _apply_rotary_emb(key_rot, cos, sin,
-                                              self.is_neox_style)
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
@@ -134,21 +139,19 @@ class MRotaryEmbedding(RotaryEmbedding):
 
         image_grid_thw = [] if image_grid_thw is None else image_grid_thw
         video_grid_thw = [] if video_grid_thw is None else video_grid_thw
-        second_per_grid_ts = [] if second_per_grid_ts is None else \
-            second_per_grid_ts
+        second_per_grid_ts = [] if second_per_grid_ts is None else second_per_grid_ts
 
-        llm_positions, mrope_position_delta = \
-            cls.get_input_positions_tensor(
-                input_tokens=input_tokens,
-                hf_config=hf_config,
-                image_grid_thw=image_grid_thw,
-                video_grid_thw=video_grid_thw,
-                second_per_grid_ts=second_per_grid_ts,
-                context_len=context_len,
-                seq_len=seq_len,
-                audio_feature_lengths=audio_feature_lengths,
-                use_audio_in_video=use_audio_in_video,
-            )
+        llm_positions, mrope_position_delta = cls.get_input_positions_tensor(
+            input_tokens=input_tokens,
+            hf_config=hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            context_len=context_len,
+            seq_len=seq_len,
+            audio_feature_lengths=audio_feature_lengths,
+            use_audio_in_video=use_audio_in_video,
+        )
 
         return llm_positions.tolist(), mrope_position_delta
 
@@ -166,6 +169,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         use_audio_in_video: bool = False,
     ) -> tuple[torch.Tensor, int]:
         from vllm.transformers_utils.config import thinker_uses_mrope
+
         if thinker_uses_mrope(hf_config):
             return cls._omni_get_input_positions_tensor(
                 input_tokens=input_tokens,
@@ -237,7 +241,8 @@ class MRotaryEmbedding(RotaryEmbedding):
 
             input_type_group: list[tuple[str, int, int]] = []
             for key, group_iter in itertools.groupby(
-                    enumerate(input_token_type), lambda x: x[1]):
+                enumerate(input_token_type), lambda x: x[1]
+            ):
                 group_list = list(group_iter)
                 start_index = group_list[0][0]
                 end_index = group_list[-1][0] + 1
@@ -246,25 +251,42 @@ class MRotaryEmbedding(RotaryEmbedding):
             video_frame_num = 1
             mm_data_idx = 0
             for modality_type, start_idx, end_idx in input_type_group:
-                st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                    llm_pos_ids_list) > 0 else 0
+                st_idx = (
+                    llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                )
                 if modality_type == "image":
                     t, h, w = (
                         image_grid_thw[mm_data_idx][0],
                         image_grid_thw[mm_data_idx][1],
                         image_grid_thw[mm_data_idx][2],
                     )
-                    llm_grid_t, llm_grid_h, llm_grid_w = \
-                        t, h // spatial_merge_size, w // spatial_merge_size
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        w // spatial_merge_size,
+                    )
 
-                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
-                        -1, llm_grid_h * llm_grid_w).flatten()
-                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                        llm_grid_t, -1, llm_grid_w).flatten()
-                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                        llm_grid_t, llm_grid_h, -1).flatten()
+                    t_index = (
+                        torch.arange(llm_grid_t)
+                        .view(-1, 1)
+                        .expand(-1, llm_grid_h * llm_grid_w)
+                        .flatten()
+                    )
+                    h_index = (
+                        torch.arange(llm_grid_h)
+                        .view(1, -1, 1)
+                        .expand(llm_grid_t, -1, llm_grid_w)
+                        .flatten()
+                    )
+                    w_index = (
+                        torch.arange(llm_grid_w)
+                        .view(1, 1, -1)
+                        .expand(llm_grid_t, llm_grid_h, -1)
+                        .flatten()
+                    )
                     llm_pos_ids_list.append(
-                        torch.stack([t_index, h_index, w_index]) + st_idx)
+                        torch.stack([t_index, h_index, w_index]) + st_idx
+                    )
                     mm_data_idx += 1
 
                 elif modality_type == "video":
@@ -273,18 +295,34 @@ class MRotaryEmbedding(RotaryEmbedding):
                         image_grid_thw[mm_data_idx][1],
                         image_grid_thw[mm_data_idx][2],
                     )
-                    llm_grid_t, llm_grid_h, llm_grid_w = \
-                        t, h // spatial_merge_size, w // spatial_merge_size
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t,
+                        h // spatial_merge_size,
+                        w // spatial_merge_size,
+                    )
 
                     for t_idx in range(llm_grid_t):
-                        t_index = torch.tensor(t_idx).view(-1, 1).expand(
-                            -1, llm_grid_h * llm_grid_w).flatten()
-                        h_index = torch.arange(llm_grid_h).view(
-                            1, -1, 1).expand(1, -1, llm_grid_w).flatten()
-                        w_index = torch.arange(llm_grid_w).view(
-                            1, 1, -1).expand(1, llm_grid_h, -1).flatten()
+                        t_index = (
+                            torch.tensor(t_idx)
+                            .view(-1, 1)
+                            .expand(-1, llm_grid_h * llm_grid_w)
+                            .flatten()
+                        )
+                        h_index = (
+                            torch.arange(llm_grid_h)
+                            .view(1, -1, 1)
+                            .expand(1, -1, llm_grid_w)
+                            .flatten()
+                        )
+                        w_index = (
+                            torch.arange(llm_grid_w)
+                            .view(1, 1, -1)
+                            .expand(1, llm_grid_h, -1)
+                            .flatten()
+                        )
                         llm_pos_ids_list.append(
-                            torch.stack([t_index, h_index, w_index]) + st_idx)
+                            torch.stack([t_index, h_index, w_index]) + st_idx
+                        )
 
                     mm_data_idx += 1
                     video_frame_num += 1
@@ -292,19 +330,17 @@ class MRotaryEmbedding(RotaryEmbedding):
                 else:
                     text_len = end_idx - start_idx
                     llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
+                        torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+                    )
                     video_frame_num = 1
 
         else:
             text_len = len(input_tokens)
-            llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1))
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1))
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
         llm_positions = llm_positions[:, context_len:seq_len]
-        mrope_position_delta = (llm_positions.max() + 1 -
-                                len(input_tokens)).item()
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
         return llm_positions, mrope_position_delta
 
     @classmethod
@@ -324,12 +360,12 @@ class MRotaryEmbedding(RotaryEmbedding):
         video_token_id = hf_config.video_token_id
         vision_start_token_id = hf_config.vision_start_token_id
         spatial_merge_size = hf_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(hf_config.vision_config,
-                                    "tokens_per_second", 1.0)
+        tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
 
         input_tokens_tensor = torch.tensor(input_tokens)
         vision_start_indices = torch.argwhere(
-            input_tokens_tensor == vision_start_token_id).squeeze(1)
+            input_tokens_tensor == vision_start_token_id
+        ).squeeze(1)
         vision_tokens = input_tokens_tensor[vision_start_indices + 1]
         image_nums = (vision_tokens == image_token_id).sum()
         video_nums = (vision_tokens == video_token_id).sum()
@@ -371,37 +407,56 @@ class MRotaryEmbedding(RotaryEmbedding):
                 remain_videos -= 1
                 ed = ed_video
 
-            llm_grid_t, llm_grid_h, llm_grid_w = \
-                t, h // spatial_merge_size, w // spatial_merge_size
+            llm_grid_t, llm_grid_h, llm_grid_w = (
+                t,
+                h // spatial_merge_size,
+                w // spatial_merge_size,
+            )
             text_len = ed - st
 
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
-            t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
-                -1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
-                       tokens_per_second).long().flatten()
+            t_index = (
+                (
+                    torch.arange(llm_grid_t)
+                    .view(-1, 1)
+                    .expand(-1, llm_grid_h * llm_grid_w)
+                    * video_second_per_grid_t
+                    * tokens_per_second
+                )
+                .long()
+                .flatten()
+            )
 
-            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                llm_grid_t, -1, llm_grid_w).flatten()
-            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                llm_grid_t, llm_grid_h, -1).flatten()
+            h_index = (
+                torch.arange(llm_grid_h)
+                .view(1, -1, 1)
+                .expand(llm_grid_t, -1, llm_grid_w)
+                .flatten()
+            )
+            w_index = (
+                torch.arange(llm_grid_w)
+                .view(1, 1, -1)
+                .expand(llm_grid_t, llm_grid_h, -1)
+                .flatten()
+            )
             llm_pos_ids_list.append(
-                torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                torch.stack([t_index, h_index, w_index]) + text_len + st_idx
+            )
             st = ed + llm_grid_t * llm_grid_h * llm_grid_w
 
         if st < len(input_tokens):
-            st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
             text_len = len(input_tokens) - st
             llm_pos_ids_list.append(
-                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+                torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+            )
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-        mrope_position_delta = (llm_positions.max() + 1 -
-                                len(input_tokens)).item()
+        mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
         llm_positions = llm_positions[:, context_len:seq_len]
 
         return llm_positions, mrope_position_delta
@@ -448,8 +503,9 @@ class MRotaryEmbedding(RotaryEmbedding):
         vision_end_token_id = thinker_config.vision_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config,
-                                    "tokens_per_second", 25)
+        tokens_per_second = getattr(
+            thinker_config.vision_config, "tokens_per_second", 25
+        )
 
         if isinstance(image_grid_thw, list):
             image_grid_thw = torch.tensor(image_grid_thw)
@@ -469,28 +525,30 @@ class MRotaryEmbedding(RotaryEmbedding):
         idx = 0
         while idx < len(src_item):
             new_src_item_len = len(new_src_item)
-            start_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                llm_pos_ids_list) > 0 else 0
-            if src_item[idx] not in [
-                    audio_token_id, video_token_id, image_token_id
-            ]:
+            start_idx = (
+                llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            )
+            if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
                 if use_audio_in_video and idx > 0:
-                    if src_item[idx] == vision_end_token_id and \
-                        src_item[idx - 1] == audio_end_token_id:
+                    if (
+                        src_item[idx] == vision_end_token_id
+                        and src_item[idx - 1] == audio_end_token_id
+                    ):
                         # processing the <|audio_eos|> before <|vision_eos|>
                         start_idx -= 1
-                    elif src_item[idx] == audio_start_token_id and \
-                        src_item[idx - 1] == vision_start_token_id:
+                    elif (
+                        src_item[idx] == audio_start_token_id
+                        and src_item[idx - 1] == vision_start_token_id
+                    ):
                         # processing the <|audio_bos|> after <|vision_eos|>
                         start_idx -= 1
                 new_src_item.append(src_item[idx])
-                llm_pos_ids = torch.tensor([start_idx],
-                                           dtype=torch.long).expand(3, -1)
+                llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
                 llm_pos_ids_list.append(llm_pos_ids)
             elif src_item[idx] == audio_token_id:
                 assert audio_seqlens is not None
                 audio_seqlen = audio_seqlens[audio_idx]
-                place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
+                place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
                 new_src_item.extend([audio_token_id] * place_num)
                 llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
                 llm_pos_ids_list.append(llm_pos_ids)
@@ -501,26 +559,30 @@ class MRotaryEmbedding(RotaryEmbedding):
                 grid_ws = image_grid_thw[:, 2]
                 t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
                 llm_pos_ids = cls._get_llm_pos_ids_for_vision(
-                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
-                    grid_ws)
+                    start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
                 llm_pos_ids_list.append(llm_pos_ids)
                 vision_seqlen = image_grid_thw[image_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 new_src_item.extend([image_token_id] * vision_seqlen)
                 image_idx += 1
             elif src_item[idx] == video_token_id and not use_audio_in_video:
                 grid_t = video_grid_thw[video_idx][0]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
-                t_index = (torch.arange(grid_t) *
-                           second_per_grid_ts[video_idx] *
-                           tokens_per_second).long()
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
                 llm_pos_ids = cls._get_llm_pos_ids_for_vision(
-                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
-                    grid_ws)
+                    start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+                )
                 llm_pos_ids_list.append(llm_pos_ids)
                 vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 new_src_item.extend([video_token_id] * vision_seqlen)
                 video_idx += 1
             else:
@@ -528,56 +590,73 @@ class MRotaryEmbedding(RotaryEmbedding):
                 assert audio_seqlens is not None
                 audio_seqlen = audio_seqlens[audio_idx]
                 vision_seqlen = video_grid_thw[video_idx].prod() // (
-                    spatial_merge_size**2)
+                    spatial_merge_size**2
+                )
                 grid_t = video_grid_thw[video_idx][0]
                 grid_h = video_grid_thw[video_idx][1]
                 grid_w = video_grid_thw[video_idx][2]
                 grid_hs = video_grid_thw[:, 1]
                 grid_ws = video_grid_thw[:, 2]
                 t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-                t_index = (torch.arange(grid_t) *
-                           second_per_grid_ts[video_idx] *
-                           tokens_per_second).long()
+                t_index = (
+                    torch.arange(grid_t)
+                    * second_per_grid_ts[video_idx]
+                    * tokens_per_second
+                ).long()
                 t_index_split_chunk = cls._split_list_into_ranges(
-                    t_index, t_ntoken_per_chunk)
+                    t_index, t_ntoken_per_chunk
+                )
                 place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
                 pure_audio_len = place_num - 2
                 added_audio_len = 0
                 audio_llm_pos_ids_list: list[torch.Tensor] = []
                 for t_chunk in t_index_split_chunk:
-                    vision_ntoken_per_chunk = len(
-                        t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
-                    new_src_item.extend([video_token_id] *
-                                        vision_ntoken_per_chunk)
+                    vision_ntoken_per_chunk = (
+                        len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+                    )
+                    new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
                     vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
-                        start_idx, video_idx, spatial_merge_size, t_chunk,
-                        grid_hs, grid_ws).split(1, dim=1)
+                        start_idx,
+                        video_idx,
+                        spatial_merge_size,
+                        t_chunk,
+                        grid_hs,
+                        grid_ws,
+                    ).split(1, dim=1)
                     llm_pos_ids_list.extend(vision_llm_pos_ids_list)
                     new_src_item.extend(
-                        min(t_ntoken_per_chunk, pure_audio_len -
-                            added_audio_len) * [audio_token_id])
-                    audio_start_idx = start_idx if len(
-                        audio_llm_pos_ids_list
-                    ) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
-                    if min(t_ntoken_per_chunk,
-                           pure_audio_len - added_audio_len) > 0:
-                        audio_llm_pos_ids_list = (torch.arange(
-                            min(t_ntoken_per_chunk, pure_audio_len -
-                                added_audio_len)).expand(3, -1) +
-                                                  audio_start_idx).split(1,
-                                                                         dim=1)
+                        min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
+                        * [audio_token_id]
+                    )
+                    audio_start_idx = (
+                        start_idx
+                        if len(audio_llm_pos_ids_list) == 0
+                        else audio_llm_pos_ids_list[-1][0].item() + 1
+                    )
+                    if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
+                        audio_llm_pos_ids_list = (
+                            torch.arange(
+                                min(
+                                    t_ntoken_per_chunk, pure_audio_len - added_audio_len
+                                )
+                            ).expand(3, -1)
+                            + audio_start_idx
+                        ).split(1, dim=1)
                     else:
                         audio_llm_pos_ids_list = []
-                    added_audio_len += min(t_ntoken_per_chunk,
-                                           pure_audio_len - added_audio_len)
+                    added_audio_len += min(
+                        t_ntoken_per_chunk, pure_audio_len - added_audio_len
+                    )
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 if added_audio_len < pure_audio_len:
                     new_src_item.extend(
-                        (pure_audio_len - added_audio_len) * [audio_token_id])
+                        (pure_audio_len - added_audio_len) * [audio_token_id]
+                    )
                     audio_llm_pos_ids_list = (
-                        torch.arange(pure_audio_len - added_audio_len).expand(
-                            3, -1) + llm_pos_ids_list[-1].max() + 1).split(
-                                1, dim=1)
+                        torch.arange(pure_audio_len - added_audio_len).expand(3, -1)
+                        + llm_pos_ids_list[-1].max()
+                        + 1
+                    ).split(1, dim=1)
                     llm_pos_ids_list.extend(audio_llm_pos_ids_list)
                 audio_idx += 1
                 video_idx += 1
@@ -585,8 +664,9 @@ class MRotaryEmbedding(RotaryEmbedding):
             idx += len(new_src_item) - new_src_item_len
 
         llm_positions = torch.cat(llm_pos_ids_list, dim=1)
-        mrope_position_delta = torch.cat(llm_pos_ids_list,
-                                         dim=1).max() + 1 - len(src_item)
+        mrope_position_delta = (
+            torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
+        )
         llm_positions = llm_positions[:, context_len:seq_len]
 
         return llm_positions, mrope_position_delta
@@ -603,22 +683,34 @@ class MRotaryEmbedding(RotaryEmbedding):
         llm_pos_ids_list = []
         llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
         llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
-        h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand(
-            len(t_index), -1, llm_grid_w).flatten())
-        w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
-            len(t_index), llm_grid_h, -1).flatten())
-        t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
-            -1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
+        h_index = (
+            torch.arange(llm_grid_h)
+            .view(1, -1, 1)
+            .expand(len(t_index), -1, llm_grid_w)
+            .flatten()
+        )
+        w_index = (
+            torch.arange(llm_grid_w)
+            .view(1, 1, -1)
+            .expand(len(t_index), llm_grid_h, -1)
+            .flatten()
+        )
+        t_index_tensor = (
+            torch.Tensor(t_index)
+            .to(llm_grid_h.device)
+            .view(-1, 1)
+            .expand(-1, llm_grid_h * llm_grid_w)
+            .long()
+            .flatten()
+        )
         _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
         llm_pos_ids_list.append(_llm_pos_ids + start_idx)
         llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
         return llm_pos_ids
 
     @staticmethod
-    def _split_list_into_ranges(lst: torch.Tensor,
-                                interval: int) -> list[list[int]]:
-        ranges: list[list[int]] = [[]
-                                   for _ in range((max(lst) // interval) + 1)]
+    def _split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
+        ranges: list[list[int]] = [[] for _ in range((max(lst) // interval) + 1)]
         for num in lst:
             index = num // interval
             ranges[index].append(num)
@@ -632,19 +724,28 @@ class MRotaryEmbedding(RotaryEmbedding):
     ) -> list[list[int]]:
         return [
             list(
-                range(context_len + mrope_position_delta,
-                      seq_len + mrope_position_delta)) for _ in range(3)
+                range(
+                    context_len + mrope_position_delta, seq_len + mrope_position_delta
+                )
+            )
+            for _ in range(3)
         ]
 
     @staticmethod
-    def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
-                                        mrope_position_delta: int,
-                                        context_len: int, num_new_tokens: int):
+    def get_next_input_positions_tensor(
+        out: np.ndarray,
+        out_offset: int,
+        mrope_position_delta: int,
+        context_len: int,
+        num_new_tokens: int,
+    ):
 
-        values = np.arange(mrope_position_delta + context_len,
-                           mrope_position_delta + context_len + num_new_tokens,
-                           dtype=out.dtype)
-        out[:, out_offset:out_offset + num_new_tokens] = values
+        values = np.arange(
+            mrope_position_delta + context_len,
+            mrope_position_delta + context_len + num_new_tokens,
+            dtype=out.dtype,
+        )
+        out[:, out_offset : out_offset + num_new_tokens] = values
 
     @classmethod
     def omni_get_updates_use_audio_in_video(
@@ -669,27 +770,28 @@ class MRotaryEmbedding(RotaryEmbedding):
         audio_end_token_id = thinker_config.audio_end_token_id
         seconds_per_chunk = thinker_config.seconds_per_chunk
         spatial_merge_size = thinker_config.vision_config.spatial_merge_size
-        tokens_per_second = getattr(thinker_config.vision_config,
-                                    "tokens_per_second", 25)
+        tokens_per_second = getattr(
+            thinker_config.vision_config, "tokens_per_second", 25
+        )
 
         grid_t = video_grid_thw[0]
         grid_h = video_grid_thw[1]
         grid_w = video_grid_thw[2]
         t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
-        t_index = (torch.arange(grid_t) * video_second_per_grid_t *
-                   tokens_per_second).long()
-        t_index_split_chunk = cls._split_list_into_ranges(
-            t_index, t_ntoken_per_chunk)
+        t_index = (
+            torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second
+        ).long()
+        t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)
 
         updates = [audio_start_token_id]
         added_audio_len = 0
         for t_chunk in t_index_split_chunk:
-            vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
-                spatial_merge_size**2)
+            vision_ntoken_per_chunk = (
+                len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
+            )
             updates.extend([video_token_id] * vision_ntoken_per_chunk)
 
-            audio_chunk_size = min(t_ntoken_per_chunk,
-                                   audio_len - added_audio_len)
+            audio_chunk_size = min(t_ntoken_per_chunk, audio_len - added_audio_len)
             updates.extend(audio_chunk_size * [audio_token_id])
             added_audio_len += audio_chunk_size
         if added_audio_len < audio_len:

vllm_omni/model_executor/models/qwen2_5_omni.py

@@ -1,77 +1,78 @@
-from functools import cached_property
-from typing import Iterable, List, Optional, Set, Tuple, Union, NamedTuple, Dict
-
-import os
 import glob
+import os
+from functools import cached_property
+from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, Union
+
 import numpy as np
 import torch
 import torch.nn as nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig, Qwen2_5OmniTalkerConfig)
-
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniTalkerConfig,
+    Qwen2_5OmniThinkerConfig,
+)
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
-    Qwen2_5OmniThinkerDummyInputsBuilder)
+)
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+
 # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model,
-                    maybe_prefix)
-from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 
 TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
 
+
 class OmniOutput(NamedTuple):
     """Output from the merged Omni model containing both text and audio."""
+
     text_hidden_states: torch.Tensor
     multimodal_outputs: dict = {}
     intermediate_tensors: Optional[IntermediateTensors] = None
 
+
 logger = init_logger(__name__)
 
+
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen2_5OmniThinkerMultiModalProcessor,
     info=Qwen2_5OmniThinkerProcessingInfo,
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                         SupportsPP,
-                                         Qwen2_5OmniConditionalGenerationMixin):
+class Qwen2_5OmniForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        self.have_multimodal_outputs = True 
+        self.have_multimodal_outputs = True
         config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         # keep vllm_config for later submodule init
         self.vllm_config = vllm_config
-        
+
         # Initialize thinker components
         thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config
         self.thinker_config = thinker_config
         self.multimodal_config = multimodal_config
-        
+
         # Initialize talker components
         talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
         self.talker_config = talker_config
-        
-
 
         self.model_stage = vllm_config.model_config.model_stage
-        if self.model_stage=="thinker":
+        if self.model_stage == "thinker":
             # Initialize thinker model (multimodal processing)
             self.thinker = init_vllm_registered_model(
                 vllm_config=vllm_config,
@@ -79,12 +80,12 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 hf_config=thinker_config,
                 # Use registry architecture key
                 architectures=["Qwen2_5OmniThinkerModel"],
-                )
+            )
             self.model = self.thinker
             self.talker = None
             self.token2wav = None
-            
-        elif self.model_stage=="talker":
+
+        elif self.model_stage == "talker":
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
             self.talker = init_vllm_registered_model(
@@ -95,14 +96,14 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 architectures=["Qwen2_5OmniTalkerModel"],
             )
             self.talker.init_multi_modal(thinker_config)
-            self.model=self.talker
+            self.model = self.talker
             self.token2wav = None
-            
-        elif self.model_stage=="code2wav":
+
+        elif self.model_stage == "code2wav":
             self.thinker = None
             self.talker = None
             # Initialize token2wav (code->mel->wav) like thinker/talker
-            self.token2wav_config = getattr(config, 'token2wav_config', None)
+            self.token2wav_config = getattr(config, "token2wav_config", None)
             self.token2wav = None
             if self.token2wav_config is not None:
                 self.token2wav = init_vllm_registered_model(
@@ -117,10 +118,13 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             self.model = self.token2wav
         else:
             raise ValueError("Invalid model stage")
-        
+
         # Set up intermediate tensors
         self.make_empty_intermediate_tensors = (
-            self.thinker.make_empty_intermediate_tensors) if self.model_stage=="thinker" else lambda: None
+            (self.thinker.make_empty_intermediate_tensors)
+            if self.model_stage == "thinker"
+            else lambda: None
+        )
 
     # -------------------- Device utilities --------------------
     @staticmethod
@@ -168,15 +172,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         multimodal_embeddings=None,
     ) -> torch.Tensor:
         if self.model_stage == "code2wav":
-            return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
-        return self.model.get_input_embeddings(
-            input_ids, multimodal_embeddings)
-       
+            return (
+                torch.zeros_like(input_ids)
+                .reshape(-1, 1)
+                .repeat(1, self.vllm_config.model_config.get_hidden_size())
+            )
+        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
 
     def get_multimodal_embeddings(self, **kwargs):
         # Delegate to thinker model for multimodal processing
         return self.model.get_multimodal_embeddings(**kwargs)
-    
+
     def last_index_of(self, list, value):
         return len(list) - 1 - list[::-1].index(value)
 
@@ -191,7 +197,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         codec: Optional[torch.Tensor] = None,
         sampling_metadata: Optional[SamplingMetadata] = None,
         logits_index: Optional[int] = None,
-        sampler = None,
+        sampler=None,
         additional_information: Optional[dict[str, object]] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
@@ -201,8 +207,8 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         2) If audio requested and codec not provided, use talker to derive codec.
         3) If audio requested (or codec provided), use token2wav to synthesize waveform.
         4) Return text hidden states (and audio when applicable).
-        """        
-        if self.model_stage=="thinker":
+        """
+        if self.model_stage == "thinker":
             # Normalize to batched inputs if caller provides 1D/2D unbatched tensors
             added_batch_dim = False
             if input_ids is not None and input_ids.ndim == 1:
@@ -215,10 +221,15 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 inputs_embeds = inputs_embeds.unsqueeze(0)
                 added_batch_dim = True
             thinker_dev = self._module_device(self.thinker)
-            
-            #if input_ids is None, set it to a zero tensor, in the length of the same as the embedding seq length
+
+            # if input_ids is None, set it to a zero tensor, in the length of the
+            # same as the embedding seq length
             if input_ids is None:
-                input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(0) #(1, 0)
+                input_ids = torch.zeros(
+                    inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev
+                ).unsqueeze(
+                    0
+                )  # (1, 0)
                 added_batch_dim = True
 
             # 1) Thinker (ensure inputs on thinker's device)
@@ -244,35 +255,59 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
 
             # Text-only path
             return OmniOutput(
-                text_hidden_states=text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states,
-                multimodal_outputs=None
+                text_hidden_states=(
+                    text_hidden_states.squeeze(0)
+                    if added_batch_dim
+                    else text_hidden_states
+                ),
+                multimodal_outputs=None,
             )
 
         # 2) Talker (if codec not provided)
-        if self.model_stage=="talker":
-            
+        if self.model_stage == "talker":
+
             if input_ids is None and additional_information is None:
-                input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
+                input_ids = torch.zeros(
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
                 additional_information = {}
                 self.thinker_reply_part = torch.zeros_like(inputs_embeds)
                 is_profile = True
             else:
                 is_profile = False
 
-            if input_ids is not None and additional_information is not None and not is_profile:
+            if (
+                input_ids is not None
+                and additional_information is not None
+                and not is_profile
+            ):
                 # read from additional_information dict
                 thinker_result = None
-                if additional_information is not None and isinstance(additional_information, dict):
+                if additional_information is not None and isinstance(
+                    additional_information, dict
+                ):
                     thinker_result = additional_information.get("thinker_result")
                     prompt_embeds = additional_information.get("prompt_embeds")
                     prompt_token_ids = additional_information.get("prompt_token_ids")
-                    thinker_output_token_ids = additional_information.get("thinker_output_token_ids")
+                    thinker_output_token_ids = additional_information.get(
+                        "thinker_output_token_ids"
+                    )
                 else:
                     thinker_result = torch.zeros_like(inputs_embeds)
                     prompt_embeds = torch.zeros_like(inputs_embeds)
-                    prompt_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
-                    thinker_output_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
-                    
+                    prompt_token_ids = torch.zeros(
+                        inputs_embeds.shape[0],
+                        dtype=torch.int64,
+                        device=inputs_embeds.device,
+                    )
+                    thinker_output_token_ids = torch.zeros(
+                        inputs_embeds.shape[0],
+                        dtype=torch.int64,
+                        device=inputs_embeds.device,
+                    )
+
                 if thinker_result is None:
                     thinker_result = torch.zeros_like(inputs_embeds)
                 self.thinker_reply_part = thinker_result.squeeze(0)
@@ -281,19 +316,24 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 input_ids, inputs_embeds = self._thinker_to_talker_prefill(
                     voice_type=voice_type,
                     output_prompt_embeds=thinker_result,
-                    output_token_ids = thinker_output_token_ids,
-                    thinker_prompt_embeds=prompt_embeds, 
-                    prompt_token_ids = prompt_token_ids,
+                    output_token_ids=thinker_output_token_ids,
+                    thinker_prompt_embeds=prompt_embeds,
+                    prompt_token_ids=prompt_token_ids,
                 )
             elif not is_profile:
-                input_ids, inputs_embeds = \
-                self._thinker_to_talker_decode_one_step(
-                    output_prompt_embeds=self.thinker_reply_part[:1] if self.thinker_reply_part.shape[0]>=1 \
-                        else torch.zeros(1, self.thinker_reply_part.shape[1]).to(self._module_device(self.model)).to(torch.bfloat16)+(-1.25*2**(-123)),
+                input_ids, inputs_embeds = self._thinker_to_talker_decode_one_step(
+                    output_prompt_embeds=(
+                        self.thinker_reply_part[:1]
+                        if self.thinker_reply_part.shape[0] >= 1
+                        else torch.zeros(1, self.thinker_reply_part.shape[1])
+                        .to(self._module_device(self.model))
+                        .to(torch.bfloat16)
+                        + (-1.25 * 2 ** (-123))
+                    ),
                     output_token_ids=input_ids,
                 )
 
-                if self.thinker_reply_part.shape[0] >=1:
+                if self.thinker_reply_part.shape[0] >= 1:
                     self.thinker_reply_part = self.thinker_reply_part[1:, :]
 
             with torch.inference_mode():
@@ -301,34 +341,49 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                     input_ids=input_ids,
                     positions=positions[0],
                     inputs_embeds=inputs_embeds,
-
                 )
-            
-            return OmniOutput(
-                text_hidden_states=talker_hidden,
-                multimodal_outputs=None
-            )
 
-        if self.model_stage=="code2wav":
-            code = input_ids if input_ids is not None else torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
-            audio_tensor = self.generate_audio(code[:-1] if code[-1]==TALKER_CODEC_EOS_TOKEN_ID else code, voice_type)
+            return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=None)
+
+        if self.model_stage == "code2wav":
+            code = (
+                input_ids
+                if input_ids is not None
+                else torch.zeros(
+                    inputs_embeds.shape[0],
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
+            )
+            audio_tensor = self.generate_audio(
+                code[:-1] if code[-1] == TALKER_CODEC_EOS_TOKEN_ID else code, voice_type
+            )
             return OmniOutput(
-                text_hidden_states = None,
-                multimodal_outputs = {
-                    "audio": audio_tensor
-                }
+                text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
             )
 
         return OmniOutput(
             text_hidden_states=torch.cat(
                 [
-                    torch.zeros([inputs_embeds.shape[0],self.talker.config.hidden_size], dtype=torch.bfloat16).to(self._module_device(self.model)),
-                    self.talker.thinker_to_talker_proj(self.talker.get_input_embeddings(torch.tensor([TALKER_CODEC_BOS_TOKEN_ID,TALKER_CODEC_EOS_TOKEN_ID]).to(torch.bfloat16).to(self._module_device(self.model))))[0]
+                    torch.zeros(
+                        [inputs_embeds.shape[0], self.talker.config.hidden_size],
+                        dtype=torch.bfloat16,
+                    ).to(self._module_device(self.model)),
+                    self.talker.thinker_to_talker_proj(
+                        self.talker.get_input_embeddings(
+                            torch.tensor(
+                                [TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]
+                            )
+                            .to(torch.bfloat16)
+                            .to(self._module_device(self.model))
+                        )
+                    )[0],
                 ],
-                dim=0),
-            multimodal_outputs=None
+                dim=0,
+            ),
+            multimodal_outputs=None,
         )
-    
+
     def generate_audio(self, code, voice_type):
         token2wav_dev = self._module_device(self.token2wav)
         if isinstance(code, torch.Tensor):
@@ -343,7 +398,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         return audio_tensor
 
     def _load_talker_embedding(
-            self,
+        self,
     ) -> torch.nn.Embedding:
         return self.talker.language_model.model.embed_tokens
 
@@ -356,19 +411,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         # embed_text_bos_token
         self.tts_text_spk_token_ids = {
             # M02: Male voice with standard Mandarin and a slight northern accent
-            'm02': 151870,
-            'Ethan': 151870,
-
+            "m02": 151870,
+            "Ethan": 151870,
             # F030: Your anime-styled virtual girlfriend
-            'f030': 151872,
-            'Chelsie': 151872,
+            "f030": 151872,
+            "Chelsie": 151872,
         }
-        self.default_tts_text_spk_type = list(
-            self.tts_text_spk_token_ids.keys())[0]
-        self.tts_text_spk_token_ids['prefix_caching'] = 151870
+        self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0]
+        self.tts_text_spk_token_ids["prefix_caching"] = 151870
 
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         self.embed_text_bos_token = self.thinker_embedding(
@@ -376,15 +429,16 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 [talker_hf_config.tts_text_start_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_text_spk_tokens = {
-            key:
-            self.thinker_embedding(
+            key: self.thinker_embedding(
                 torch.tensor(
                     [value],
                     dtype=torch.long,
                     device=self._module_device(self.talker),
-                ))
+                )
+            )
             for key, value in self.tts_text_spk_token_ids.items()
         }
         self.embed_text_eos_token = self.thinker_embedding(
@@ -392,31 +446,36 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 [talker_hf_config.tts_text_end_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_text_pad_token = self.thinker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_text_pad_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_bos_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_start_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_eos_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_end_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         self.embed_codec_pad_token = self.talker_embedding(
             torch.tensor(
                 [talker_hf_config.tts_codec_pad_token_id],
                 dtype=torch.long,
                 device=self._module_device(self.talker),
-            ))
+            )
+        )
         return set(["thinker_embedding.weight", "talker_embedding.weight"])
 
     def _get_embed_text_spk_token(self, voice_type: str):
@@ -426,7 +485,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
 
     def _get_text_spk_token_id(self, voice_type: str):
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         if voice_type not in self.tts_text_spk_token_ids:
@@ -439,52 +498,58 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         output_prompt_embeds,
         output_token_ids,
         thinker_prompt_embeds,
-        prompt_token_ids
+        prompt_token_ids,
     ):
 
         talker_hf_config = self.talker_config
-        if hasattr(talker_hf_config, 'talker_config'):
+        if hasattr(talker_hf_config, "talker_config"):
             talker_hf_config = talker_hf_config.talker_config
 
         # if len(output.outputs[0].token_ids) == 2:
-            # issue request
-        prompt_embeds = torch.cat([
-            thinker_prompt_embeds,
-            self._get_embed_text_spk_token(voice_type) +
-            self.embed_codec_pad_token,
-            output_prompt_embeds[:1] + self.embed_codec_bos_token,
-        ],
-                                    dim=0)
+        # issue request
+        prompt_embeds = torch.cat(
+            [
+                thinker_prompt_embeds,
+                self._get_embed_text_spk_token(voice_type) + self.embed_codec_pad_token,
+                output_prompt_embeds[:1] + self.embed_codec_bos_token,
+            ],
+            dim=0,
+        )
 
         prompt_token_ids_processed = prompt_token_ids + [
             talker_hf_config.tts_codec_pad_token_id,
             output_token_ids[0],
-
         ]
         input_tokens_len = len(prompt_token_ids_processed)
         # the code below is from model runner in Qwen, may need to further discuss later
         if input_tokens_len > 2:
-            prompt_token_ids_processed = (
-                [self.talker_config.tts_codec_mask_token_id] *
-                (input_tokens_len - 2) + [
-                    self.talker_config.tts_codec_pad_token_id,
-                    self.talker_config.tts_codec_start_token_id
-                ])
+            prompt_token_ids_processed = [
+                self.talker_config.tts_codec_mask_token_id
+            ] * (input_tokens_len - 2) + [
+                self.talker_config.tts_codec_pad_token_id,
+                self.talker_config.tts_codec_start_token_id,
+            ]
         else:
             prompt_token_ids_processed = [
                 self.talker_config.tts_codec_pad_token_id,
                 self.talker_config.tts_codec_start_token_id,
             ][-input_tokens_len:]
-        if isinstance(prompt_token_ids_processed,list):
-            prompt_token_ids_processed = torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
+        if isinstance(prompt_token_ids_processed, list):
+            prompt_token_ids_processed = (
+                torch.Tensor(prompt_token_ids_processed)
+                .to(torch.int64)
+                .to(self._module_device(self.talker))
+            )
         return prompt_token_ids_processed, prompt_embeds
-            
+
     def _thinker_to_talker_decode_one_step(
         self,
         output_prompt_embeds,
         output_token_ids,
     ):
-        processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids) #for decode
+        processed_output_token_embeds = (
+            output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids)
+        )  # for decode
         return output_token_ids, processed_output_token_embeds
 
     def compute_logits(
@@ -495,7 +560,7 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         # Handle OmniOutput type
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
-        
+
         # Use thinker model for logits computation
         return self.model.compute_logits(hidden_states, sampling_metadata)
 
@@ -511,29 +576,28 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         """
         Generate speech from text tokens using the talker and token2wav models.
         This method is kept for backward compatibility and direct speech generation.
-        
+
         Args:
             text_tokens: Text tokens from thinker model
             voice_type: Voice type for speech generation
-            
+
         Returns:
             Audio tensor
         """
         # Generate codec tokens using talker model
         talker_output = self.talker(
-            input_ids=None,
-            positions=None,
-            inputs_embeds=text_tokens
+            input_ids=None, positions=None, inputs_embeds=text_tokens
         )
-        
+
         # Convert talker output to codec tokens
         codec_tokens = self._convert_to_codec_tokens(talker_output)
-        
+
         # Generate audio using token2wav model
         return self._codec_to_audio(codec_tokens, voice_type=voice_type)
 
-
-    def _convert_to_codec_tokens(self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor:
+    def _convert_to_codec_tokens(
+        self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata
+    ) -> torch.Tensor:
         """
         Reference (HF): use the talker's codec head to obtain logits, suppress BOS,
         then greedily select the next codec token for the current step.
@@ -541,12 +605,19 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         with torch.inference_mode():
             logits = self.talker.compute_logits(talker_output, None)
             if logits is None:
-                return torch.zeros((talker_output.size(0), 0), dtype=torch.long, device=talker_output.device)
+                return torch.zeros(
+                    (talker_output.size(0), 0),
+                    dtype=torch.long,
+                    device=talker_output.device,
+                )
 
-            # Suppress only codec_bos, consistent with HF generate's suppress_tokens behavior
+            # Suppress only codec_bos, consistent with HF generate's
+            # suppress_tokens behavior
             bos_id = None
-            if hasattr(self, 'talker_config') and hasattr(self.talker_config, 'tts_codec_start_token_id'):
-                bos_id = int(getattr(self.talker_config, 'tts_codec_start_token_id'))
+            if hasattr(self, "talker_config") and hasattr(
+                self.talker_config, "tts_codec_start_token_id"
+            ):
+                bos_id = int(getattr(self.talker_config, "tts_codec_start_token_id"))
             if bos_id is not None:
                 logits[..., bos_id] = -1e9
 
@@ -555,20 +626,25 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             return next_id.to(dtype=torch.long)
 
     def _init_token2wav_model(self):
-        """Initialize speaker resources if provided; model is constructed in __init__."""
+        """Initialize speaker resources if provided; model is constructed in
+        __init__."""
         if self.token2wav is None or self.token2wav_config is None:
             return
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         # optional speaker resources
-        conds = getattr(self.token2wav_config, 'conds', None)
-        ref_mels = getattr(self.token2wav_config, 'ref_mels', None)
+        conds = getattr(self.token2wav_config, "conds", None)
+        ref_mels = getattr(self.token2wav_config, "ref_mels", None)
         if isinstance(conds, dict) and isinstance(ref_mels, dict):
-            self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
-            self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
+            self._token2wav_conds = {
+                k: torch.as_tensor(v, device=device) for k, v in conds.items()
+            }
+            self._token2wav_ref_mels = {
+                k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()
+            }
         # legacy: load from directory if provided
         model_path = self.vllm_config.model_config.model
         if isinstance(model_path, str) and os.path.isdir(model_path):
-            spk_pt = os.path.join(model_path, 'spk_dict.pt')
+            spk_pt = os.path.join(model_path, "spk_dict.pt")
             if os.path.exists(spk_pt):
                 data = torch.load(spk_pt, map_location=device)
                 for key, value in data.items():
@@ -576,20 +652,30 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                     self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
             else:
                 # legacy npy inputs
-                for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*spk_emb.npy'))):
-                    key = os.path.basename(f).split('_')[0].lower()
-                    self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
-                for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*ref_mel.npy'))):
-                    key = os.path.basename(f).split('_')[0].lower()
-                    self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*spk_emb.npy"))
+                ):
+                    key = os.path.basename(f).split("_")[0].lower()
+                    self._token2wav_conds[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
+                for f in sorted(
+                    glob.glob(os.path.join(model_path, "inputs", "*ref_mel.npy"))
+                ):
+                    key = os.path.basename(f).split("_")[0].lower()
+                    self._token2wav_ref_mels[key] = torch.as_tensor(
+                        np.load(f), device=device
+                    )
 
-    def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
+    def _codec_to_audio(
+        self, codec_tokens: torch.Tensor, voice_type: str = "default"
+    ) -> Optional[torch.Tensor]:
         if self.token2wav is None:
             self._init_token2wav_model()
         if self.token2wav is None:
             return None
         # Normalize voice type
-        voice = (voice_type or 'default')
+        voice = voice_type or "default"
         # Resolve cond / ref_mel if provided
         cond = None
         ref_mel = None
@@ -599,9 +685,17 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         # Fallback: create dummy cond/ref_mel if not provided
         token2wav_dev = self._module_device(self.token2wav)
         if cond is None:
-            cond = torch.zeros((1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32)
+            cond = torch.zeros(
+                (1, self.token2wav_config.dit_config.enc_emb_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
         if ref_mel is None:
-            ref_mel = torch.zeros((1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32)
+            ref_mel = torch.zeros(
+                (1, 300, self.token2wav_config.dit_config.mel_dim),
+                device=token2wav_dev,
+                dtype=torch.float32,
+            )
 
         # Ensure codec is (1, T) long tensor on correct device
         if isinstance(codec_tokens, torch.Tensor):
@@ -609,37 +703,49 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             if codec.ndim == 1:
                 codec = codec.unsqueeze(0)
         else:
-            codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
+            codec = torch.as_tensor(
+                codec_tokens, dtype=torch.long, device=token2wav_dev
+            ).unsqueeze(0)
 
-        # Streaming with chunked process and boundary alignment (rely on token2wav.process_chunk)
-        factor = getattr(self.token2wav.token2wav.factor, 'factor', 2)
+        # Streaming with chunked process and boundary alignment
+        # (rely on token2wav.process_chunk)
+        factor = getattr(self.token2wav.token2wav.factor, "factor", 2)
         chunk_size = 48
-        mel_dim = getattr(self.token2wav.token2wav.code2wav_dit_model, 'mel_dim', self.token2wav_config.dit_config.mel_dim)
+        mel_dim = getattr(
+            self.token2wav.token2wav.code2wav_dit_model,
+            "mel_dim",
+            self.token2wav_config.dit_config.mel_dim,
+        )
         total_mel = int(codec.shape[1] * factor)
         steps = 10
 
         # Prepare initial noise for the whole sequence
-        y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
+        y_all = torch.randn(
+            (1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev
+        )
 
-        logger.info(f"Currently, we do not use the chunked process, we only use the token2wav.process_chunk for the whole sequence.\
-                    The stream mode will be implemented in the future.")
+        logger.info(
+            "Currently, we do not use the chunked process, we only use the "
+            "token2wav.process_chunk for the whole sequence. "
+            "The stream mode will be implemented in the future."
+        )
 
         chunk_ends = []
         for i in range(codec.shape[1]):
             chunk_code_length = i * 2 - 24
-            finished = i==(codec.shape[1]-1)
-            if (chunk_code_length > 0 and
-                chunk_code_length % chunk_size == 0) or finished:
+            finished = i == (codec.shape[1] - 1)
+            if (
+                chunk_code_length > 0 and chunk_code_length % chunk_size == 0
+            ) or finished:
                 chunk_ends.append(i)
 
         # Number of chunks in mel domain
         prev_generated = None
         wav_chunks: list = []
-        prev_id = 0
 
         with torch.inference_mode():
-            for n,i in enumerate([0]):
-                finished = (i == codec.shape[1] - 1)
+            for n, i in enumerate([0]):
+                finished = i == codec.shape[1] - 1
                 _, audio_chunk = self.token2wav.process_chunk(
                     conditioning=cond,
                     reference_mel=ref_mel,
@@ -652,7 +758,6 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
                 )
                 prev_generated = audio_chunk
                 wav_chunks.append(audio_chunk.detach().cpu().numpy())
-                prev_id = i
 
         if len(wav_chunks) == 0:
             return torch.zeros(0, device=token2wav_dev)
@@ -667,11 +772,11 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
         talker_weights = []
         token2wav_weights = []
         for k, v in weights:
-            if k.startswith('thinker.'):
+            if k.startswith("thinker."):
                 thinker_weights.append((k, v))
-            elif k.startswith('talker.'):
+            elif k.startswith("talker."):
                 talker_weights.append((k, v))
-            elif k.startswith('token2wav.'):
+            elif k.startswith("token2wav."):
                 token2wav_weights.append((k, v))
             else:
                 raise ValueError(f"Unknown weight prefix: {k}")
@@ -681,30 +786,42 @@ class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
             if thinker_weights:
                 thinker_loaded = self.thinker.load_weights(thinker_weights)
             else:
-                thinker_loaded = set([k for k,v in thinker_weights])
-            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, 'thinker')
+                thinker_loaded = set([k for k, v in thinker_weights])
+            thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
             loaded_weights.update(thinker_loaded)
 
-        
         # Load talker weights
         if talker_weights and self.talker is not None:
             # Map talker weights to appropriate components
             if self.thinker is None:
-                thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
+                thinker_embedding_weights = [
+                    w
+                    for n, w in thinker_weights
+                    if n == "thinker.model.embed_tokens.weight"
+                ]
                 if thinker_embedding_weights:
-                    self.thinker_embedding = nn.Embedding(thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1])
-                    self.thinker_embedding.weight = nn.Parameter(thinker_embedding_weights[0].to(self._module_device(self.talker)))
+                    self.thinker_embedding = nn.Embedding(
+                        thinker_embedding_weights[0].shape[0],
+                        thinker_embedding_weights[0].shape[1],
+                    )
+                    self.thinker_embedding.weight = nn.Parameter(
+                        thinker_embedding_weights[0].to(
+                            self._module_device(self.talker)
+                        )
+                    )
             talker_loaded = self.talker.load_weights(talker_weights)
-            talker_loaded = add_prefix_to_loaded_weights(talker_loaded, 'talker')
+            talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
             loaded_weights.update(talker_loaded)
             loaded_weights.update(self._init_special_tokens_embeddings())
-            
-        
+
         # Load token2wav weights (if any)
         if token2wav_weights and self.token2wav is not None:
             self._init_token2wav_model()
-            t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"))
-            t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, 'token2wav')
+            t2w_loaded = self.token2wav.load_weights(
+                token2wav_weights,
+                os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"),
+            )
+            t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
             loaded_weights.update(t2w_loaded)
-        
-        return loaded_weights
+
+        return loaded_weights

vllm_omni/model_executor/models/qwen2_5_omni_talker.py

@@ -1,32 +1,42 @@
 from functools import cached_property
-from typing import Iterable, List, Optional, Set, Tuple, Union
+from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 import torch.nn as nn
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniTalkerConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniAudioEncoder,
+)
 
-from vllm.attention import AttentionMetadata
+# from vllm.attention import AttentionMetadata  # unused import
 from vllm.config import VllmConfig
-from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
     Qwen2_5OmniThinkerProcessingInfo,
-    Qwen2_5OmniThinkerDummyInputsBuilder)
-from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.logger import init_logger
-from vllm.sequence import IntermediateTensors
-
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model,
-                    maybe_prefix,
-                    merge_multimodal_embeddings)
+)
 from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
 
 
 @MULTIMODAL_REGISTRY.register_processor(
@@ -34,9 +44,9 @@ from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
     info=Qwen2_5OmniThinkerProcessingInfo,
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                                SupportsPP,
-                                                Qwen2_5OmniConditionalGenerationMixin):
+class Qwen2_5OmniTalkerForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
     logger = init_logger(__name__)
     # Align to thinker-style static mapper for clarity
     hf_to_vllm_mapper = WeightsMapper(
@@ -48,7 +58,8 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
             # fallback root
             "talker.": "",
-        })
+        }
+    )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -60,8 +71,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
 
         if hasattr(config, "talker_config"):
             self.config = config.talker_config
-            vllm_config.model_config.hf_text_config = \
-                    vllm_config.model_config.hf_config.talker_config
+            vllm_config.model_config.hf_text_config = (
+                vllm_config.model_config.hf_config.talker_config
+            )
         else:
             self.config = config
 
@@ -76,12 +88,13 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         self.language_model = init_vllm_registered_model(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "language_model"),
-            hf_config=getattr(self.config, 'text_config', self.config),
+            hf_config=getattr(self.config, "text_config", self.config),
             architectures=["Qwen2ForCausalLM_old"],
         )
         self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors)
-    
+            self.language_model.make_empty_intermediate_tensors
+        )
+
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -107,17 +120,20 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None \
-            and len(multimodal_embeddings) != 0:
+        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
 
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, multimodal_embeddings, [
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
                     self.config.image_token_index,
                     self.config.video_token_index,
-                    self.config.audio_token_index
-                ])
+                    self.config.audio_token_index,
+                ],
+            )
         return inputs_embeds
 
     def forward(
@@ -128,9 +144,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         inputs_embeds: Optional[torch.Tensor] = None,
         **kwargs: object,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
-        forward_context: ForwardContext = get_forward_context()
-        attn_metadata: AttentionMetadata = forward_context.attn_metadata
+        assert (
+            input_ids is not None or inputs_embeds is not None
+        ), "input_ids or inputs_embeds must be provided"
+        # forward_context: ForwardContext = get_forward_context()  # unused variable
 
         if intermediate_tensors is not None:
             inputs_embeds = None
@@ -143,10 +160,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         # projection
         inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds)
 
-        hidden_states = self.language_model.model(input_ids,
-                                                  positions,
-                                                  intermediate_tensors,
-                                                  inputs_embeds=inputs_embeds)
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
         return hidden_states
 
     def compute_logits(
@@ -154,8 +170,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
         hidden_states: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states,
-                                                  sampling_metadata)
+        return self.language_model.compute_logits(hidden_states, sampling_metadata)
 
     def sample(
         self,
@@ -164,11 +179,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
     ) -> Optional[SamplerOutput]:
         return self.language_model.sample(logits, sampling_metadata)
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
-            skip_prefixes=['thinker.', 'token2wav.'],
+            skip_prefixes=["thinker.", "token2wav."],
         )
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
@@ -180,42 +194,53 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             device = next(self.parameters()).device
             self.logger.info(
                 "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__, True, total_bytes / (1024**2), str(device))
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
         except Exception:
             pass
         multi_model_weights = set()
         for name, param in self.visual.named_parameters():
-            multi_model_weights.add("visual."+name)
+            multi_model_weights.add("visual." + name)
         for name, param in self.audio_tower.named_parameters():
-            multi_model_weights.add("audio_tower."+name)
+            multi_model_weights.add("audio_tower." + name)
         loaded.update(multi_model_weights)
         return loaded
-    
+
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
 
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds"
-                             ) and "image" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds"
-                             ) and "video" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features"
-                             ) and "audio" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
-    def get_multimodal_embeddings(self,
-                                  **kwargs: object) -> MultiModalEmbeddings:
+    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
 
-        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
-            **kwargs)
+        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
 
@@ -236,4 +261,4 @@ class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
             if modality == "audio":
                 audio_embeddings = self._process_audio_input(multimodal_input)
                 multimodal_embeddings += audio_embeddings
-        return multimodal_embeddings
+        return multimodal_embeddings

vllm_omni/model_executor/models/qwen2_5_omni_thinker.py

@@ -7,45 +7,73 @@ import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig)
+    Qwen2_5OmniConfig,
+    Qwen2_5OmniThinkerConfig,
+)
 from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniAudioEncoder)
+    Qwen2_5OmniAudioEncoder,
+)
 from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
-    Qwen2_5OmniProcessor)
+    Qwen2_5OmniProcessor,
+)
 from transformers.models.whisper import WhisperFeatureExtractor
-
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
-    Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs,
-    Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs,
-    Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs,
-    Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs)
+    Qwen2_5_VisionTransformer,
+    Qwen2_5_VLImageEmbeddingInputs,
+    Qwen2_5_VLImageInputs,
+    Qwen2_5_VLImagePixelInputs,
+    Qwen2_5_VLProcessingInfo,
+    Qwen2_5_VLVideoEmbeddingInputs,
+    Qwen2_5_VLVideoInputs,
+    Qwen2_5_VLVideoPixelInputs,
+)
 from vllm.model_executor.models.qwen2_audio import (
-    Qwen2AudioInputs, Qwen2AudioProcessingInfo,
-    _get_feat_extract_output_lengths)
+    Qwen2AudioInputs,
+    Qwen2AudioProcessingInfo,
+    _get_feat_extract_output_lengths,
+)
 from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    init_vllm_registered_model,
+    maybe_prefix,
+    merge_multimodal_embeddings,
+)
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (ImageItem, ModalityData,
-                                    MultiModalDataDict, MultiModalFieldConfig,
-                                    MultiModalKwargs, NestedTensors)
-from vllm.multimodal.parse import (AudioProcessorItems, DictEmbeddingItems,
-                                   ModalityDataItems, MultiModalDataItems,
-                                   MultiModalDataParser)
-from vllm.multimodal.processing import (BaseMultiModalProcessor,
-                                        PlaceholderFeaturesInfo,
-                                        PromptReplacement, PromptUpdate)
+from vllm.multimodal.inputs import (
+    ImageItem,
+    ModalityData,
+    MultiModalDataDict,
+    MultiModalFieldConfig,
+    MultiModalKwargs,
+    NestedTensors,
+)
+from vllm.multimodal.parse import (
+    AudioProcessorItems,
+    DictEmbeddingItems,
+    ModalityDataItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+from vllm.multimodal.processing import (
+    BaseMultiModalProcessor,
+    PlaceholderFeaturesInfo,
+    PromptReplacement,
+    PromptUpdate,
+)
 from vllm.multimodal.profiling import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
 
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
-                    init_vllm_registered_model, maybe_prefix,
-                    merge_multimodal_embeddings)
-
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 
 try:
@@ -57,8 +85,7 @@ logger = init_logger(__name__)
 
 
 def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
-    audio_feature_lengths = hf_inputs.get("audio_feature_lengths",
-                                          torch.empty((0, )))
+    audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
 
     image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
     image_grid_sizes = image_grid_thw.prod(-1)
@@ -68,18 +95,17 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
 
     return dict(
         input_audio_features=MultiModalFieldConfig.flat_from_sizes(
-            "audio", audio_feature_lengths, dim=1),
+            "audio", audio_feature_lengths, dim=1
+        ),
         feature_attention_mask=MultiModalFieldConfig.batched("audio"),
         audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
-        pixel_values=MultiModalFieldConfig.flat_from_sizes(
-            "image", image_grid_sizes),
-        image_embeds=MultiModalFieldConfig.flat_from_sizes(
-            "image", image_grid_sizes),
+        pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
+        image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),
         image_grid_thw=MultiModalFieldConfig.batched("image"),
         pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes),
-        video_embeds=MultiModalFieldConfig.flat_from_sizes(
-            "video", video_grid_sizes),
+            "video", video_grid_sizes
+        ),
+        video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes),
         video_grid_thw=MultiModalFieldConfig.batched("video"),
         second_per_grid_ts=MultiModalFieldConfig.batched("video"),
     )
@@ -95,17 +121,16 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
             return DictEmbeddingItems(
                 data,
                 modality="audio",
-                required_fields={
-                    "input_audio_features", "audio_feature_lengths"
-                },
+                required_fields={"input_audio_features", "audio_feature_lengths"},
                 fields_factory=_qwen2_5_omni_thinker_field_config,
             )
 
         return super()._parse_audio_data(data)
 
 
-class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
-                                       Qwen2_5_VLProcessingInfo):
+class Qwen2_5OmniThinkerProcessingInfo(
+    Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
+):
 
     def get_hf_config(self):
         return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
@@ -128,7 +153,8 @@ class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
 
 
 class Qwen2_5OmniThinkerDummyInputsBuilder(
-        BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
+    BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
+):
 
     def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
         num_audios = mm_counts.get("audio", 0)
@@ -141,8 +167,11 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
         image_token: str = hf_processor.image_token
         video_token: str = hf_processor.video_token
 
-        return (audio_token * num_audios + image_token * num_images +
-                video_token * num_videos)
+        return (
+            audio_token * num_audios
+            + image_token * num_images
+            + video_token * num_videos
+        )
 
     def get_dummy_mm_data(
         self,
@@ -155,40 +184,45 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
 
         feature_extractor = self.info.get_feature_extractor()
 
-        target_audio_length = min(
-            feature_extractor.chunk_length,
-            30,
-        ) * feature_extractor.sampling_rate
-        target_width, target_height = \
-            self.info.get_image_size_with_most_features()
-        target_num_frames = \
-            self.info.get_num_frames_with_most_features(seq_len, mm_counts)
+        target_audio_length = (
+            min(
+                feature_extractor.chunk_length,
+                30,
+            )
+            * feature_extractor.sampling_rate
+        )
+        target_width, target_height = self.info.get_image_size_with_most_features()
+        target_num_frames = self.info.get_num_frames_with_most_features(
+            seq_len, mm_counts
+        )
 
         mm_data = {
-            "audio":
-            self._get_dummy_audios(length=target_audio_length,
-                                   num_audios=num_audios),
-            "image":
-            self._get_dummy_images(width=target_width,
-                                   height=target_height,
-                                   num_images=num_images),
-            "video":
-            self._get_dummy_videos(width=target_width,
-                                   height=target_height,
-                                   num_frames=target_num_frames,
-                                   num_videos=num_videos),
+            "audio": self._get_dummy_audios(
+                length=target_audio_length, num_audios=num_audios
+            ),
+            "image": self._get_dummy_images(
+                width=target_width, height=target_height, num_images=num_images
+            ),
+            "video": self._get_dummy_videos(
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
+            ),
         }
 
         return mm_data
 
 
 class Qwen2_5OmniThinkerMultiModalProcessor(
-        BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
+    BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
+):
 
     def _get_data_parser(self) -> MultiModalDataParser:
         feature_extractor = self.info.get_feature_extractor()
         return Qwen2_5OmniThinkerMultiModalDataParser(
-            target_sr=feature_extractor.sampling_rate)
+            target_sr=feature_extractor.sampling_rate
+        )
 
     def _call_hf_processor(
         self,
@@ -204,7 +238,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         if audios:
             # NOTE: Qwen2.5-Omni processor accept "audio"
             mm_data["audio"] = audios
-            mm_kwargs = dict(**mm_kwargs, )
+            mm_kwargs = dict(
+                **mm_kwargs,
+            )
 
         hf_inputs = super()._call_hf_processor(
             prompt=prompt,
@@ -213,17 +249,19 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             tok_kwargs=tok_kwargs,
         )
 
-        input_features = hf_inputs.pop('input_features', None)
-        feature_attention_mask = hf_inputs.get('feature_attention_mask', None)
-        if ('input_audio_features' not in hf_inputs
-                and input_features is not None):
+        input_features = hf_inputs.pop("input_features", None)
+        feature_attention_mask = hf_inputs.get("feature_attention_mask", None)
+        if "input_audio_features" not in hf_inputs and input_features is not None:
             if feature_attention_mask is not None:
-                input_features = input_features.permute(
-                    0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
-            hf_inputs['input_audio_features'] = input_features
-        if ('audio_feature_lengths' not in hf_inputs
-                and feature_attention_mask is not None):
-            hf_inputs['audio_feature_lengths'] = feature_attention_mask.sum(-1)
+                input_features = input_features.permute(0, 2, 1)[
+                    feature_attention_mask.bool()
+                ].permute(1, 0)
+            hf_inputs["input_audio_features"] = input_features
+        if (
+            "audio_feature_lengths" not in hf_inputs
+            and feature_attention_mask is not None
+        ):
+            hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1)
         return hf_inputs
 
     def _get_mm_fields_config(
@@ -249,14 +287,12 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             hf_processor_mm_kwargs,
             mm_kwargs,
         )
-        mm_prompt_updates = self._bind_and_group_updates(
-            unbound_prompt_updates)
+        mm_prompt_updates = self._bind_and_group_updates(unbound_prompt_updates)
 
         mm_item_counts = mm_items.get_all_counts()
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
 
         if is_update_applied:
             mm_placeholders = self._find_mm_placeholders(
@@ -265,9 +301,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_item_counts,
             )
             self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-                use_audio_in_video=use_audio_in_video)
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+            )
 
             tokenizer = self.info.get_tokenizer()
             prompt = decode_tokens(tokenizer, prompt_ids)
@@ -282,9 +317,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_item_counts,
             )
             self._validate_mm_placeholders(
-                mm_placeholders,
-                mm_item_counts,
-                use_audio_in_video=use_audio_in_video)
+                mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
+            )
 
         tokenizer = self.info.get_tokenizer()
         prompt = decode_tokens(tokenizer, prompt_ids)
@@ -302,8 +336,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
     ) -> Sequence[PromptUpdate]:
         processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
         tokenizer = self.info.get_tokenizer()
-        image_processor = self.info.get_image_processor(
-            **hf_processor_mm_kwargs)
+        image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
         vocab = tokenizer.get_vocab()
 
         audio_token = processor.audio_token
@@ -319,12 +352,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             audio_output_lengths = []
         elif audio_feature_lengths is not None:
             _, audio_output_lens = _get_feat_extract_output_lengths(
-                audio_feature_lengths)
+                audio_feature_lengths
+            )
             audio_output_lengths = audio_output_lens.tolist()
         elif feature_attention_mask is not None:
             assert isinstance(feature_attention_mask, torch.Tensor)
             _, audio_output_lens = _get_feat_extract_output_lengths(
-                feature_attention_mask.sum(-1))
+                feature_attention_mask.sum(-1)
+            )
             audio_output_lengths = audio_output_lens.tolist()
 
         # number of audios read from video.
@@ -339,7 +374,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 audio = audios.get(item_idx)
                 raise ValueError(
                     f"The audio {audio} (len={len(audio)}) is too short "
-                    "to be represented inside the model")
+                    "to be represented inside the model"
+                )
 
             return [audio_token_id] * num_features
 
@@ -351,21 +387,20 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             token_id = image_token_id if modality == "image" else video_token_id
             return [token_id] * (int(grid_thw.prod()) // merge_length)
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
         thinker_config = self.info.get_hf_config()
 
         def get_replacement_qwen2_use_audio_in_video(item_idx: int):
             nonlocal audio_in_video_item_idx
 
-            audio_num_features = audio_output_lengths[audio_in_video_item_idx +
-                                                      item_idx]
+            audio_num_features = audio_output_lengths[
+                audio_in_video_item_idx + item_idx
+            ]
             video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx]
 
             audio_in_video_item_idx += 1
 
-            second_per_grid_ts = hf_processor_mm_kwargs.get(
-                "second_per_grid_ts", None)
+            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
             if second_per_grid_ts:
                 video_second_per_grid_t = second_per_grid_ts[item_idx]
             else:
@@ -379,8 +414,10 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             )
 
         video_replacement_fn = (
-            get_replacement_qwen2_use_audio_in_video if use_audio_in_video else
-            partial(get_replacement_qwen2_vision, modality="video"))
+            get_replacement_qwen2_use_audio_in_video
+            if use_audio_in_video
+            else partial(get_replacement_qwen2_vision, modality="video")
+        )
 
         return [
             PromptReplacement(
@@ -391,8 +428,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             PromptReplacement(
                 modality="image",
                 target=image_token,
-                replacement=partial(get_replacement_qwen2_vision,
-                                    modality="image"),
+                replacement=partial(get_replacement_qwen2_vision, modality="image"),
             ),
             PromptReplacement(
                 modality="video",
@@ -445,8 +481,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         """
         mm_counts = mm_items.get_all_counts()
 
-        use_audio_in_video = hf_processor_mm_kwargs.get(
-            "use_audio_in_video", False)
+        use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)
         if use_audio_in_video and "video" in mm_counts:
             assert "audio" in mm_counts
             mm_counts["audio"] -= mm_counts["video"]
@@ -476,36 +511,43 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 
 class Qwen2_5OmniConditionalGenerationMixin:
 
-    def _validate_and_reshape_mm_tensor(self,
-                                        mm_input: object,
-                                        name: str,
-                                        dim: int = 0) -> torch.Tensor:
+    def _validate_and_reshape_mm_tensor(
+        self, mm_input: object, name: str, dim: int = 0
+    ) -> torch.Tensor:
         if not isinstance(mm_input, (torch.Tensor, list)):
-            raise ValueError(f"Incorrect type of {name}. "
-                             f"Got type: {type(mm_input)}")
+            raise ValueError(
+                f"Incorrect type of {name}. " f"Got type: {type(mm_input)}"
+            )
         if isinstance(mm_input, torch.Tensor):
             return torch.concat(list(mm_input), dim=dim)
         else:
             return torch.concat(mm_input, dim=dim)
 
     def _parse_and_validate_audio_input(
-            self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
-        input_audio_features = kwargs.pop('input_audio_features', None)
-        audio_feature_lengths = kwargs.pop('audio_feature_lengths', None)
-        feature_attention_mask = kwargs.pop('feature_attention_mask', None)
+        self, **kwargs: object
+    ) -> Optional[Qwen2AudioInputs]:
+        input_audio_features = kwargs.pop("input_audio_features", None)
+        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
+        feature_attention_mask = kwargs.pop("feature_attention_mask", None)
         if input_audio_features is None:
             return None
         input_audio_features = self._validate_and_reshape_mm_tensor(
-            input_audio_features, 'input_audio_features', dim=1)
+            input_audio_features, "input_audio_features", dim=1
+        )
         if feature_attention_mask is not None:
             feature_attention_mask = self._validate_and_reshape_mm_tensor(
-                feature_attention_mask, 'feature_attention_mask')
+                feature_attention_mask, "feature_attention_mask"
+            )
         if not isinstance(input_audio_features, (torch.Tensor, list)):
-            raise ValueError("Incorrect type of audio input features. "
-                             f"Got type: {type(input_audio_features)}")
-        return Qwen2AudioInputs(input_features=input_audio_features,
-                                audio_feature_lengths=audio_feature_lengths,
-                                feature_attention_mask=feature_attention_mask)
+            raise ValueError(
+                "Incorrect type of audio input features. "
+                f"Got type: {type(input_audio_features)}"
+            )
+        return Qwen2AudioInputs(
+            input_features=input_audio_features,
+            audio_feature_lengths=audio_feature_lengths,
+            feature_attention_mask=feature_attention_mask,
+        )
 
     def _parse_and_validate_image_input(
         self,
@@ -520,31 +562,42 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if pixel_values is not None:
             pixel_values = self._validate_and_reshape_mm_tensor(
-                pixel_values, "image pixel values")
+                pixel_values, "image pixel values"
+            )
             image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw")
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(pixel_values, (torch.Tensor, list)):
-                raise ValueError("Incorrect type of image pixel values. "
-                                 f"Got type: {type(pixel_values)}")
+                raise ValueError(
+                    "Incorrect type of image pixel values. "
+                    f"Got type: {type(pixel_values)}"
+                )
 
-            return Qwen2_5_VLImagePixelInputs(type="pixel_values",
-                                              pixel_values=pixel_values,
-                                              image_grid_thw=image_grid_thw)
+            return Qwen2_5_VLImagePixelInputs(
+                type="pixel_values",
+                pixel_values=pixel_values,
+                image_grid_thw=image_grid_thw,
+            )
 
         if image_embeds is not None:
             image_embeds = self._validate_and_reshape_mm_tensor(
-                image_embeds, "image embeds")
+                image_embeds, "image embeds"
+            )
             image_grid_thw = self._validate_and_reshape_mm_tensor(
-                image_grid_thw, "image grid_thw")
+                image_grid_thw, "image grid_thw"
+            )
 
             if not isinstance(image_embeds, torch.Tensor):
-                raise ValueError("Incorrect type of image embeddings. "
-                                 f"Got type: {type(image_embeds)}")
+                raise ValueError(
+                    "Incorrect type of image embeddings. "
+                    f"Got type: {type(image_embeds)}"
+                )
             return Qwen2_5_VLImageEmbeddingInputs(
                 type="image_embeds",
                 image_embeds=image_embeds,
-                image_grid_thw=image_grid_thw)
+                image_grid_thw=image_grid_thw,
+            )
 
     def _parse_and_validate_video_input(
         self,
@@ -559,9 +612,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if pixel_values_videos is not None:
             pixel_values_videos = self._validate_and_reshape_mm_tensor(
-                pixel_values_videos, "video pixel values")
+                pixel_values_videos, "video pixel values"
+            )
             video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw")
+                video_grid_thw, "video grid_thw"
+            )
 
             return Qwen2_5_VLVideoPixelInputs(
                 type="pixel_values_videos",
@@ -571,17 +626,22 @@ class Qwen2_5OmniConditionalGenerationMixin:
 
         if video_embeds is not None:
             video_embeds = self._validate_and_reshape_mm_tensor(
-                video_embeds, "video embeds")
+                video_embeds, "video embeds"
+            )
             video_grid_thw = self._validate_and_reshape_mm_tensor(
-                video_grid_thw, "video grid_thw")
+                video_grid_thw, "video grid_thw"
+            )
 
             if not isinstance(video_embeds, torch.Tensor):
-                raise ValueError("Incorrect type of video embeddings. "
-                                 f"Got type: {type(video_embeds)}")
+                raise ValueError(
+                    "Incorrect type of video embeddings. "
+                    f"Got type: {type(video_embeds)}"
+                )
             return Qwen2_5_VLVideoEmbeddingInputs(
                 type="video_embeds",
                 video_embeds=video_embeds,
-                video_grid_thw=video_grid_thw)
+                video_grid_thw=video_grid_thw,
+            )
 
     def _process_audio_input(
         self,
@@ -596,16 +656,18 @@ class Qwen2_5OmniConditionalGenerationMixin:
             assert input_features.shape[0] == 1
             input_features = input_features.squeeze(0)
         if audio_feature_lengths.ndim == 2:
-            assert audio_feature_lengths.shape[
-                0] == 1 or audio_feature_lengths.shape[1] == 1
+            assert (
+                audio_feature_lengths.shape[0] == 1
+                or audio_feature_lengths.shape[1] == 1
+            )
             if audio_feature_lengths.shape[0] == 1:
                 audio_feature_lengths = audio_feature_lengths.squeeze(0)
             else:
                 audio_feature_lengths = audio_feature_lengths.squeeze(1)
 
         audio_feat_lengths, audio_output_lengths = (
-            self.audio_tower._get_feat_extract_output_lengths(
-                audio_feature_lengths))
+            self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)
+        )
 
         audio_outputs = self.audio_tower(
             input_features.to(self.audio_tower.dtype),
@@ -616,8 +678,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
         return audio_features.split(audio_output_lengths.tolist())
 
     def _process_image_input(
-            self,
-            image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
+        self, image_input: Qwen2_5_VLImageInputs
+    ) -> tuple[torch.Tensor, ...]:
         if image_input["type"] == "image_embeds":
             return image_input["image_embeds"].type(self.visual.dtype)
 
@@ -633,18 +695,18 @@ class Qwen2_5OmniConditionalGenerationMixin:
         return image_embeds.split(sizes.tolist())
 
     def _process_video_input(
-            self,
-            video_input: Qwen2_5_VLVideoInputs,
-            video_hashes: list[str] = None,
-            cached_video_embeds: torch.Tensor = None) -> torch.Tensor:
+        self,
+        video_input: Qwen2_5_VLVideoInputs,
+        video_hashes: list[str] = None,
+        cached_video_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
         if video_input["type"] == "video_embeds":
             return video_input["video_embeds"].type(self.visual.dtype)
 
         grid_thw = video_input["video_grid_thw"]
         assert grid_thw.ndim == 2
 
-        pixel_values_videos = video_input["pixel_values_videos"].type(
-            self.visual.dtype)
+        pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
         video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
         # Split concatenated embeddings for each video item.
         merge_size = self.visual.spatial_merge_size
@@ -659,14 +721,15 @@ class Qwen2_5OmniConditionalGenerationMixin:
     dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
 )
 class Qwen2_5OmniThinkerForConditionalGeneration(
-        nn.Module, SupportsMultiModal, SupportsPP,
-        Qwen2_5OmniConditionalGenerationMixin):
+    nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
+):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "thinker.lm_head.": "language_model.lm_head.",
             "thinker.model.": "language_model.model.",
             "thinker.": "",
-        })
+        }
+    )
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
@@ -681,8 +744,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        thinker_config: Qwen2_5OmniThinkerConfig = (
-            vllm_config.model_config.hf_config)
+        thinker_config: Qwen2_5OmniThinkerConfig = vllm_config.model_config.hf_config
         quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         self.config = thinker_config
@@ -698,7 +760,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             logger.warning(
                 "flash_attn is not available, the model may not yield the "
                 "exactly same result as the transformers implementation "
-                "in the audio tower part.")
+                "in the audio tower part."
+            )
 
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -716,7 +779,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         )
 
         self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors)
+            self.language_model.make_empty_intermediate_tensors
+        )
 
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}
@@ -724,28 +788,35 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # Preserve the order of modalities if there are multiple of them
         # from the order of kwargs.
         for input_key in kwargs:
-            if input_key in ("pixel_values", "image_embeds"
-                             ) and "image" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "image"] = self._parse_and_validate_image_input(**kwargs)
-            if input_key in ("pixel_values_videos", "video_embeds"
-                             ) and "video" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "video"] = self._parse_and_validate_video_input(**kwargs)
-            if input_key in ("input_audio_features"
-                             ) and "audio" not in mm_input_by_modality:
-                mm_input_by_modality[
-                    "audio"] = self._parse_and_validate_audio_input(**kwargs)
+            if (
+                input_key in ("pixel_values", "image_embeds")
+                and "image" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("pixel_values_videos", "video_embeds")
+                and "video" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
+                    **kwargs
+                )
+            if (
+                input_key in ("input_audio_features")
+                and "audio" not in mm_input_by_modality
+            ):
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
+                    **kwargs
+                )
         return mm_input_by_modality
 
     def get_language_model(self) -> torch.nn.Module:
         return self.language_model
 
-    def get_multimodal_embeddings(self,
-                                  **kwargs: object) -> MultiModalEmbeddings:
+    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
 
-        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(
-            **kwargs)
+        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not mm_input_by_modality:
             return []
 
@@ -774,21 +845,23 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
     ) -> torch.Tensor:
         inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None \
-            and len(multimodal_embeddings) != 0:
+        if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
 
             # TODO (ywang96): support overlapping modalitiy embeddings so that
             # `use_audio_in_video` will work on V1.
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, multimodal_embeddings, [
+                input_ids,
+                inputs_embeds,
+                multimodal_embeddings,
+                [
                     self.config.image_token_index,
                     self.config.video_token_index,
-                    self.config.audio_token_index
-                ])
+                    self.config.audio_token_index,
+                ],
+            )
         return inputs_embeds
 
-    def get_multimodal_embeddings_v0(
-            self, **kwargs: object) -> Optional[NestedTensors]:
+    def get_multimodal_embeddings_v0(self, **kwargs: object) -> Optional[NestedTensors]:
         audio_input = self._parse_and_validate_audio_input(**kwargs)
         image_input = self._parse_and_validate_image_input(**kwargs)
         video_input = self._parse_and_validate_video_input(**kwargs)
@@ -826,7 +899,8 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             if modality == "video":
                 placeholder_token_id = self.config.video_token_index
             inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, embeddings, placeholder_token_id)
+                input_ids, inputs_embeds, embeddings, placeholder_token_id
+            )
         return inputs_embeds
 
     def forward(
@@ -844,35 +918,40 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
         # condition is for v0 compatibility.
         elif inputs_embeds is None:
             multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
-            inputs_embeds = self.get_input_embeddings(
-                input_ids, multimodal_embeddings)
-            text_inputs_embeds = self.get_input_embeddings(input_ids, [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] if multimodal_embeddings is not None else None)
+            inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings)
+            text_inputs_embeds = self.get_input_embeddings(
+                input_ids,
+                (
+                    [
+                        (torch.zeros_like(embeddings), modality)
+                        for embeddings, modality in multimodal_embeddings
+                    ]
+                    if multimodal_embeddings is not None
+                    else None
+                ),
+            )
             input_ids = None
         else:
             text_inputs_embeds = inputs_embeds
 
-        hidden_states = self.language_model.model(input_ids,
-                                                  positions,
-                                                  intermediate_tensors,
-                                                  inputs_embeds=inputs_embeds)
-        return text_inputs_embeds, hidden_states.unsqueeze(0) #(1, S, D)
+        hidden_states = self.language_model.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
+        )
+        return text_inputs_embeds, hidden_states.unsqueeze(0)  # (1, S, D)
 
     def compute_logits(
         self,
         hidden_states: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> Optional[torch.Tensor]:
-        return self.language_model.compute_logits(hidden_states,
-                                                  sampling_metadata)
+        return self.language_model.compute_logits(hidden_states, sampling_metadata)
 
-    def load_weights(self, weights: Iterable[tuple[str,
-                                                   torch.Tensor]]) -> set[str]:
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,
             skip_prefixes=["talker.", "token2wav."],
         )
-        loaded_weights = loader.load_weights(weights,
-                                             mapper=self.hf_to_vllm_mapper)
+        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
         try:
             total_bytes = 0
@@ -882,8 +961,11 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             device = next(self.parameters()).device
             logger.info(
                 "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__, True, total_bytes / (1024**2), str(device))
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
         except Exception:
             pass
         return loaded_weights
-

vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py

@@ -2,50 +2,53 @@
 #      Start Token2Wav     #
 ############################
 
-from typing import Iterable, Optional, Set, Tuple, List, Union
+import math
+from typing import Iterable, List, Optional, Set, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import numpy as np
-import math
 from torch.nn import Parameter
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniBigVGANConfig,
+    Qwen2_5OmniDiTConfig,
+    Qwen2_5OmniToken2WavConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniPreTrainedModel,
+)
 
+# Bring in HF base classes, configs and utilities used below
+from transformers.utils.logging import get_logger as _hf_get_logger
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.sequence import IntermediateTensors
-from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
 from vllm.model_executor.models.interfaces import SupportsPP
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader as _Vllm_AutoWeightsLoader,
-    WeightsMapper as _Vllm_WeightsMapper,
+)
+from vllm.model_executor.models.utils import WeightsMapper as _Vllm_WeightsMapper
+from vllm.model_executor.models.utils import (
     init_vllm_registered_model as _Vllm_init_vllm_registered_model,
-    maybe_prefix as _Vllm_maybe_prefix,
 )
+from vllm.model_executor.models.utils import maybe_prefix as _Vllm_maybe_prefix
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
 
-# Bring in HF base classes, configs and utilities used below
-from transformers.utils import ModelOutput
-from transformers.utils.logging import get_logger as _hf_get_logger
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
-from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
-    Qwen2_5OmniPreTrainedModel,
-    Qwen2_5OmniPreTrainedModelForConditionalGeneration,
-)
-from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
-    Qwen2_5OmniDiTConfig,
-    Qwen2_5OmniBigVGANConfig,
-    Qwen2_5OmniToken2WavConfig,
-)
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
 def auto_docstring(func=None, **_kwargs):
     if func is None:
+
         def wrapper(f):
             return f
+
         return wrapper
     return func
 
+
 # HF logger alias
 logger = _hf_get_logger(__name__)
 
@@ -183,7 +186,7 @@ class AttentiveStatisticsPooling(nn.Module):
     def _length_to_mask(self, length, max_len=None, dtype=None, device=None):
         """Creates a binary mask for each sequence.
 
-        Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
+        Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3  # noqa: E501
 
         Arguments
         ---------
@@ -213,7 +216,9 @@ class AttentiveStatisticsPooling(nn.Module):
 
     def _compute_statistics(self, x, m, dim=2):
         mean = (m * x).sum(dim)
-        std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
+        std = torch.sqrt(
+            (m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps)
+        )
         return mean, std
 
     def forward(self, hidden_states):
@@ -222,7 +227,10 @@ class AttentiveStatisticsPooling(nn.Module):
 
         # Make binary mask of shape [N, 1, L]
         mask = self._length_to_mask(
-            lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device
+            lengths * seq_length,
+            max_len=seq_length,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
         )
         mask = mask.unsqueeze(1)
 
@@ -272,7 +280,9 @@ class SqueezeExcitationRes2NetBlock(nn.Module):
             kernel_size=1,
             dilation=1,
         )
-        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
+        self.res2net_block = Res2NetBlock(
+            out_channels, out_channels, res2net_scale, kernel_size, dilation
+        )
         self.tdnn2 = TimeDelayNetBlock(
             out_channels,
             out_channels,
@@ -300,10 +310,13 @@ class ECAPA_TimeDelayNet(torch.nn.Module):
 
     def __init__(self, config: Qwen2_5OmniDiTConfig):
         super().__init__()
-        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
-            config.enc_dilations
-        ):
-            raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
+        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(
+            config.enc_channels
+        ) != len(config.enc_dilations):
+            raise ValueError(
+                "enc_channels, enc_kernel_sizes and enc_dilations should have "
+                "same length"
+            )
         self.channels = config.enc_channels
         self.blocks = nn.ModuleList()
 
@@ -397,14 +410,26 @@ class DiTInputEmbedding(nn.Module):
     ):
         if apply_cfg:
             hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
-            speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
-            condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
+            speaker_embedding = torch.cat(
+                [speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0
+            )
+            condition_vector = torch.cat(
+                [condition_vector, torch.zeros_like(condition_vector)], dim=0
+            )
             code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
         elif drop_audio_cond:  # cfg for cond audio
             condition_vector = torch.zeros_like(condition_vector)
             speaker_embedding = torch.zeros_like(speaker_embedding)
-        condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
-        hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
+        condition_vector = (
+            self.spk_encoder(condition_vector)
+            .unsqueeze(1)
+            .repeat(1, hidden_states.size(1), 1)
+        )
+        hidden_states = self.proj(
+            torch.cat(
+                (hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1
+            )
+        )
 
         return hidden_states
 
@@ -438,9 +463,13 @@ class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
 
     def forward(self, hidden_states, emb=None):
         emb = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(
+            emb, 6, dim=1
+        )
 
-        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        hidden_states = (
+            self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        )
         return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
 
@@ -459,7 +488,9 @@ class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
         emb = self.linear(self.silu(emb))
         scale, shift = torch.chunk(emb, 2, dim=1)
 
-        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        hidden_states = (
+            self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+        )
         return hidden_states
 
 
@@ -496,14 +527,18 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
         position_ids (`torch.Tensor`, *optional*):
             Deprecated and unused.
         unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+            The 'unsqueeze_dim' argument specifies the dimension along which to
+            unsqueeze cos[position_ids] and sin[position_ids] so that they can be
+            properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape
+            [batch_size, seq_len, head_dim]. Then, if q and k have the shape
+            [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1
+            makes cos[position_ids] and sin[position_ids] broadcastable to the
+            shapes of q and k. Similarly, if q and k have the shape
+            [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
     Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated
+        using the Rotary Position Embedding.
     """
 
     def rotate_half_codec(x):
@@ -535,7 +570,9 @@ class DiTAttention(nn.Module):
         self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
         self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
 
-        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
+        self.to_out = nn.ModuleList(
+            [nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)]
+        )
 
     def forward(
         self,
@@ -558,9 +595,12 @@ class DiTAttention(nn.Module):
         value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
 
         # apply rotary position embedding
-        # Due to training process, only first head is applied with RoPE, will be fixed at next release
+        # Due to training process, only first head is applied with RoPE,
+        # will be fixed at next release
         cos, sin = position_embeddings
-        query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
+        query[:, :1], key[:, :1] = apply_rotary_pos_emb(
+            query[:, :1], key[:, :1], cos, sin
+        )
 
         attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
         attention_weights, _ = attention_interface(
@@ -572,8 +612,11 @@ class DiTAttention(nn.Module):
             is_causal=False,
         )
 
-        # mask. e.g. inference got a batch with different target durations, mask out the padding
-        attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
+        # mask. e.g. inference got a batch with different target durations,
+        # mask out the padding
+        attention_weights = attention_weights.reshape(
+            batch_size, -1, self.heads * head_dim
+        )
         attention_weights = attention_weights.to(query.dtype)
 
         # linear proj
@@ -603,7 +646,9 @@ class DiTTimestepEmbedding(nn.Module):
     def __init__(self, dim, freq_embed_dim=256):
         super().__init__()
         self.time_embed = SinusPositionEmbedding(freq_embed_dim)
-        self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
+        self.time_mlp = nn.ModuleList(
+            [nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]
+        )
 
     def forward(self, timestep):  # noqa: F821
         time_hidden = self.time_embed(timestep)
@@ -614,21 +659,29 @@ class DiTTimestepEmbedding(nn.Module):
 
 
 class DiTDecoderLayer(nn.Module):
-    def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
+    def __init__(
+        self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0
+    ):
         super().__init__()
         self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
 
         self.attn = DiTAttention(config)
         self.look_ahead_block = look_ahead_block
         self.look_backward_block = look_backward_block
-        self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
+        self.ff_norm = nn.LayerNorm(
+            config.hidden_size, elementwise_affine=False, eps=1e-6
+        )
+        self.ff = DiTMLP(
+            dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout
+        )
 
     def forward(
         self, hidden_states, timestep, position_embeddings=None, block_diff=None
     ):  # x: noised input, t: time embedding
         # pre-norm & modulation for attention input
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(
+            hidden_states, emb=timestep
+        )
 
         # attention
         attn_output = self.attn(
@@ -641,7 +694,9 @@ class DiTDecoderLayer(nn.Module):
         # process attention output for input x
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
 
-        norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        norm = (
+            self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        )
         ff_output = self.ff(norm)
         hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
 
@@ -650,7 +705,8 @@ class DiTDecoderLayer(nn.Module):
 
 class SnakeBeta(nn.Module):
     """
-    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    A modified Snake function which uses separate parameters for the magnitude
+    of the periodic components
     Shape:
         - Input: (B, C, T)
         - Output: (B, C, T), same shape as the input
@@ -658,7 +714,8 @@ class SnakeBeta(nn.Module):
         - alpha - trainable parameter that controls frequency
         - beta - trainable parameter that controls magnitude
     References:
-        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        - This activation function is a modified version based on this paper
+          by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
         https://huggingface.co/papers/2006.08195
     """
 
@@ -682,9 +739,9 @@ class SnakeBeta(nn.Module):
         beta = self.beta.unsqueeze(0).unsqueeze(-1)
         alpha = torch.exp(alpha)
         beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
+        hidden_states = hidden_states + (
+            1.0 / (beta + self.no_div_by_zero)
+        ) * torch.pow(torch.sin(hidden_states * alpha), 2)
 
         return hidden_states
 
@@ -714,7 +771,9 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
     else:
         beta = 0.0
 
-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+    kaiser_window = torch.kaiser_window(
+        kernel_size, beta=beta, periodic=False, dtype=torch.float32
+    )
 
     # Compute time indices
     if is_even:
@@ -724,7 +783,9 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
 
     # Compute sinc filter
     if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
+        return torch.zeros(
+            (1, 1, kernel_size), dtype=torch.float32
+        )  # Ensures correct shape
 
     sinc_filter = torch.sinc(2 * cutoff * time_indices)
     normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -739,21 +800,32 @@ class UpSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
         self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
         self.stride = ratio
         self.pad = self.kernel_size // ratio - 1
         self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
 
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
         self.register_buffer("filter", filter, persistent=False)
 
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
+        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(
+            self.filter.dtype
+        )
         hidden_states = self.ratio * F.conv_transpose1d(
-            hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
+            hidden_states,
+            self.filter.expand(channels, -1, -1),
+            stride=self.stride,
+            groups=channels,
         ).to(hidden_states_dtype)
         hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
 
@@ -781,8 +853,15 @@ class DownSample1d(nn.Module):
     def forward(self, hidden_states):
         channels = hidden_states.shape[1]
         hidden_states_dtype = hidden_states.dtype
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
-        out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels).to(hidden_states_dtype)
+        hidden_states = F.pad(
+            hidden_states, (self.pad_left, self.pad_right), mode="replicate"
+        ).to(self.filter.dtype)
+        out = F.conv1d(
+            hidden_states,
+            self.filter.expand(channels, -1, -1),
+            stride=self.stride,
+            groups=channels,
+        ).to(hidden_states_dtype)
         return out
 
 
@@ -877,10 +956,15 @@ class AMPBlock(torch.nn.Module):
             ]
         )
 
-        self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # total number of conv layers
 
         self.activations = nn.ModuleList(
-            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+            [
+                TorchActivation1d(activation=SnakeBeta(channels))
+                for _ in range(self.num_layers)
+            ]
         )
 
     def _get_padding(self, kernel_size, dilation=1):
@@ -901,7 +985,8 @@ class AMPBlock(torch.nn.Module):
 
 @auto_docstring(
     custom_intro="""
-    The full Qwen2.5Omni Token2WavBigVGAN model. Which take mel spectrogram as input and predict waveform.
+    The full Qwen2.5Omni Token2WavBigVGAN model. Which take mel spectrogram
+    as input and predict waveform.
     """
 )
 class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
@@ -912,7 +997,9 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         self.num_residual_blocks = len(config.resblock_kernel_sizes)
         self.num_upsample_layers = len(config.upsample_rates)
 
-        self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
+        self.conv_pre = nn.Conv1d(
+            config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3
+        )
 
         # Removing extra ModuleList breaks official state dict
         ups = [
@@ -927,31 +1014,54 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
                     )
                 ]
             )
-            for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
+            for layer_idx, (stride, kernel_size) in enumerate(
+                zip(config.upsample_rates, config.upsample_kernel_sizes)
+            )
         ]
         self.ups = nn.ModuleList(ups)
 
         self.resblocks = nn.ModuleList(
             [
-                AMPBlock(config.upsample_initial_channel // (2 ** (layer_idx + 1)), kernel_size, dilation)
+                AMPBlock(
+                    config.upsample_initial_channel // (2 ** (layer_idx + 1)),
+                    kernel_size,
+                    dilation,
+                )
                 for layer_idx in range(self.num_upsample_layers)
-                for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
+                for kernel_size, dilation in zip(
+                    config.resblock_kernel_sizes, config.resblock_dilation_sizes
+                )
             ]
         )
 
         self.activation_post = TorchActivation1d(
-            activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
+            activation=SnakeBeta(
+                config.upsample_initial_channel // (2**self.num_upsample_layers)
+            )
         )
         self.conv_post = nn.Conv1d(
-            config.upsample_initial_channel // (2**self.num_upsample_layers), 1, 7, 1, padding=3, bias=False
+            config.upsample_initial_channel // (2**self.num_upsample_layers),
+            1,
+            7,
+            1,
+            padding=3,
+            bias=False,
         )
 
     def normalize_spectrogram(self, spectrogram, max_value, min_db):
-        return torch.clamp((2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value, -max_value, max_value)
+        return torch.clamp(
+            (2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value,
+            -max_value,
+            max_value,
+        )
 
     def amplitude_to_db(self, amplitude, min_db_level):
         min_level = torch.exp(
-            torch.tensor(min_db_level / 20.0 * np.log(10), device=amplitude.device, dtype=amplitude.dtype)
+            torch.tensor(
+                min_db_level / 20.0 * np.log(10),
+                device=amplitude.device,
+                dtype=amplitude.dtype,
+            )
         )
         return 20 * torch.log10(torch.clamp(amplitude, min=min_level))
 
@@ -967,7 +1077,9 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
         for layer_index in range(self.num_upsample_layers):
             hidden_representation = self.ups[layer_index][0](hidden_representation)
             residual_output = sum(
-                self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
+                self.resblocks[layer_index * self.num_residual_blocks + block_index](
+                    hidden_representation
+                )
                 for block_index in range(self.num_residual_blocks)
             )
             residual_output = residual_output / self.num_residual_blocks
@@ -986,20 +1098,48 @@ class RungeKutta4ODESolver:
         self._one_third = 1 / 3
         self._two_thirds = 2 / 3
 
-    def _rk4_step(self, function, time_start, time_step, time_end, value_start, function_value_start=None):
-        k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
-        k2 = function(time_start + time_step * self._one_third, value_start + time_step * k1 * self._one_third)
-        k3 = function(time_start + time_step * self._two_thirds, value_start + time_step * (k2 - k1 * self._one_third))
+    def _rk4_step(
+        self,
+        function,
+        time_start,
+        time_step,
+        time_end,
+        value_start,
+        function_value_start=None,
+    ):
+        k1 = (
+            function_value_start
+            if function_value_start is not None
+            else function(time_start, value_start)
+        )
+        k2 = function(
+            time_start + time_step * self._one_third,
+            value_start + time_step * k1 * self._one_third,
+        )
+        k3 = function(
+            time_start + time_step * self._two_thirds,
+            value_start + time_step * (k2 - k1 * self._one_third),
+        )
         k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
         return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
 
     def _compute_step(self, function, time_start, time_step, time_end, value_start):
         function_value_start = function(time_start, value_start)
-        return self._rk4_step(
-            function, time_start, time_step, time_end, value_start, function_value_start=function_value_start
-        ), function_value_start
+        return (
+            self._rk4_step(
+                function,
+                time_start,
+                time_step,
+                time_end,
+                value_start,
+                function_value_start=function_value_start,
+            ),
+            function_value_start,
+        )
 
-    def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+    def _linear_interpolation(
+        self, time_start, time_end, value_start, value_end, time_point
+    ):
         if time_point == time_start:
             return value_start
         if time_point == time_end:
@@ -1020,12 +1160,21 @@ class RungeKutta4ODESolver:
         current_value = self.initial_value
         for time_start, time_end in zip(time_points[:-1], time_points[1:]):
             time_step = time_end - time_start
-            delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
+            delta_value, _ = self._compute_step(
+                self.function, time_start, time_step, time_end, current_value
+            )
             next_value = current_value + delta_value
 
-            while current_index < len(time_points) and time_end >= time_points[current_index]:
+            while (
+                current_index < len(time_points)
+                and time_end >= time_points[current_index]
+            ):
                 solution[current_index] = self._linear_interpolation(
-                    time_start, time_end, current_value, next_value, time_points[current_index]
+                    time_start,
+                    time_end,
+                    current_value,
+                    next_value,
+                    time_points[current_index],
                 )
                 current_index += 1
 
@@ -1036,7 +1185,8 @@ class RungeKutta4ODESolver:
 
 @auto_docstring(
     custom_intro="""
-    The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as input and predict mel spectrogram.
+    The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as
+    input and predict mel spectrogram.
     """
 )
 class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
@@ -1049,7 +1199,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         self.repeats = config.repeats
         self.time_embed = DiTTimestepEmbedding(config.hidden_size)
 
-        self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
+        self.text_embed = DiTCodecEmbedding(
+            config.num_embeds, config.emb_dim, config.repeats
+        )
         self.input_embed = DiTInputEmbedding(config)
 
         self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config.head_dim)
@@ -1069,12 +1221,16 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 )
             )
 
-        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size)  # final modulation
+        self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(
+            config.hidden_size
+        )  # final modulation
         self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
 
     def _create_block_diff(self, hidden_states):
         batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
-        block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size  # [seq_length]
+        block_indices = (
+            torch.arange(seq_len, device=hidden_states.device) // self.block_size
+        )  # [seq_length]
 
         block_i = block_indices.unsqueeze(1)  # [seq_length, 1]
         block_j = block_indices.unsqueeze(0)  # [1, seq_length]
@@ -1099,8 +1255,12 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
 
         # Compute embeddings
         time_embedding = self.time_embed(time_step)
-        text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
-        text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+        text_embedding = self.text_embed(
+            quantized_code, drop_code=False if apply_cfg else drop_code
+        )
+        text_embedding_unconditioned = (
+            self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+        )
 
         hidden_states = self.input_embed(
             hidden_states,
@@ -1140,11 +1300,17 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         guidance_scale=0.5,
         sway_coefficient=-1.0,
     ):
-        noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
+        noise_initialization = torch.randn(
+            [1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype
+        )
         maximum_duration = quantized_code.shape[1] * self.repeats
-        initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
+        initial_state = noise_initialization[:, :maximum_duration].to(
+            quantized_code.device
+        )
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
+            1, maximum_duration, 1
+        )
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1171,17 +1337,28 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+            return (
+                guided_prediction
+                + (guided_prediction - null_prediction) * guidance_scale
+            )
 
         initial_time = 0
         time_embedding = torch.linspace(
-            initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype
+            initial_time,
+            1,
+            num_steps,
+            device=quantized_code.device,
+            dtype=conditioning_vector.dtype,
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+            time_embedding += sway_coefficient * (
+                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
+            )
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(
+            function=ode_function, initial_value=initial_state
+        )
         solution_trajectory = ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
@@ -1212,7 +1389,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         """
         initial_state = y0.to(quantized_code.device)
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, initial_state.shape[1], 1)
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
+            1, initial_state.shape[1], 1
+        )
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1239,17 +1418,28 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                 apply_cfg=True,
             )
             guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+            return (
+                guided_prediction
+                + (guided_prediction - null_prediction) * guidance_scale
+            )
 
         initial_time = 0
         time_embedding = torch.linspace(
-            initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype
+            initial_time,
+            1,
+            num_steps,
+            device=quantized_code.device,
+            dtype=conditioning_vector.dtype,
         )
 
         if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+            time_embedding += sway_coefficient * (
+                torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
+            )
 
-        ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+        ode_solver = RungeKutta4ODESolver(
+            function=ode_function, initial_value=initial_state
+        )
         solution_trajectory = ode_solver.integrate(time_embedding)
 
         generated_waveform = solution_trajectory[-1]
@@ -1259,26 +1449,34 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
 
 @auto_docstring(
     custom_intro="""
-    The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform.
+    The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech
+    tokens as input and predict mel spectrogram and a BigVGAN vocoder take
+    mel spectrogram as input and predict waveform.
     """
 )
 class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
     config: Qwen2_5OmniToken2WavConfig
     base_model_prefix = "model"
-    _no_split_modules = ["Qwen2_5OmniToken2WavDiTModel", "Qwen2_5OmniToken2WavBigVGANModel"]
+    _no_split_modules = [
+        "Qwen2_5OmniToken2WavDiTModel",
+        "Qwen2_5OmniToken2WavBigVGANModel",
+    ]
 
     def __init__(self, config: Qwen2_5OmniToken2WavConfig):
         super().__init__(config)
         attn_impl = config._attn_implementation
         if config._attn_implementation == "flash_attention_2":
             logger.warning_once(
-                "Qwen2_5OmniToken2WavModel must inference with fp32, but flash_attention_2 only supports fp16 and bf16, "
-                "attention implementation of Qwen2_5OmniToken2WavModel will fallback to sdpa."
+                "Qwen2_5OmniToken2WavModel must inference with fp32, but "
+                "flash_attention_2 only supports fp16 and bf16, "
+                "attention implementation of Qwen2_5OmniToken2WavModel will "
+                "fallback to sdpa."
             )
             attn_impl = "sdpa"
         elif config._attn_implementation == "eager":
             logger.warning_once(
-                "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
+                "Qwen2_5OmniToken2WavModel does not support eager attention "
+                "implementation, fall back to sdpa"
             )
             attn_impl = "sdpa"
         self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
@@ -1301,7 +1499,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
 
         # codec embedding size for masking EOS out-of-range
         try:
-            self.codec_embed_size = self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
+            self.codec_embed_size = (
+                self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
+            )
         except Exception:
             self.codec_embed_size = -1
 
@@ -1341,7 +1541,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
 
         return waveform
 
-    # ============== Chunked processing helpers (compat with qwen2_code2wav_dit) ==============
+    # ============== Chunked processing helpers (compat with qwen2_code2wav_dit) ==============  # noqa: E501
     @torch.inference_mode()
     def process_chunk_dit_batch(
         self,
@@ -1351,7 +1551,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         y0: torch.Tensor,
         steps: int,
     ) -> torch.Tensor:
-        """Block-wise DiT: generate mel from initial state y0 for the given code slice."""
+        """
+        Block-wise DiT: generate mel from initial state y0 for the given code slice.
+        """
         # prevent codec out-of-range (eos)
         if self.codec_embed_size > 0:
             code = code.clone()
@@ -1384,11 +1586,21 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
     ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
         """Streaming per small chunk: returns (mel_or_None, audio_slice)."""
         start_index = max(i * self.chunk_size - self.past_cache_size, 0)
-        end_index = min((i + 1) * self.chunk_size + self.future_cache_size,
-                        codec_all.shape[1] * self.factor)
+        end_index = min(
+            (i + 1) * self.chunk_size + self.future_cache_size,
+            codec_all.shape[1] * self.factor,
+        )
 
-        y0 = y_all[:, start_index:end_index].reshape(1, -1, self.code2wav_dit_model.mel_dim).contiguous()
-        codec = codec_all[:, start_index // self.factor:end_index // self.factor].reshape(1, -1).contiguous()
+        y0 = (
+            y_all[:, start_index:end_index]
+            .reshape(1, -1, self.code2wav_dit_model.mel_dim)
+            .contiguous()
+        )
+        codec = (
+            codec_all[:, start_index // self.factor : end_index // self.factor]
+            .reshape(1, -1)
+            .contiguous()
+        )
 
         # generate mel for current window (B, mel_dim, T)
         generated = self.process_chunk_dit_batch(
@@ -1459,29 +1671,42 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         # Normalize dtype
         generated = generated.to(torch.float32)
         if i == 0:
-            mel = generated[:, :, :self.chunk_size]
+            mel = generated[:, :, : self.chunk_size]
         elif finished:
-            mel_trim = generated[:, :, self.past_cache_size:]
-            mel = torch.cat([prev_generated[:, :, -self.future_size * 2:], mel_trim], dim=2)
+            mel_trim = generated[:, :, self.past_cache_size :]
+            mel = torch.cat(
+                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
+            )
         else:
             if start_index == 0:
-                mel_trim = generated[:, :, i * self.chunk_size:-self.future_cache_size]
+                mel_trim = generated[
+                    :, :, i * self.chunk_size : -self.future_cache_size
+                ]
             else:
-                mel_trim = generated[:, :, self.past_cache_size:-self.future_cache_size]
-            mel = torch.cat([prev_generated[:, :, -self.future_size * 2:], mel_trim], dim=2)
+                mel_trim = generated[
+                    :, :, self.past_cache_size : -self.future_cache_size
+                ]
+            mel = torch.cat(
+                [prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
+            )
 
         audio = self.code2wav_bigvgan_model(mel)
         if i == 0:
-            audio_output = audio[:-self.future_size * self.vocoder_hop]
+            audio_output = audio[: -self.future_size * self.vocoder_hop]
         elif finished:
-            audio_output = audio[self.future_size * self.vocoder_hop:]
+            audio_output = audio[self.future_size * self.vocoder_hop :]
         else:
-            audio_output = audio[self.future_size * self.vocoder_hop:-self.future_size * self.vocoder_hop]
+            audio_output = audio[
+                self.future_size
+                * self.vocoder_hop : -self.future_size
+                * self.vocoder_hop
+            ]
         return mel, audio_output
 
 
 # ================= vLLM-style wrapper for Token2Wav =================
 
+
 class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     logger = init_logger(__name__)
 
@@ -1490,7 +1715,8 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         orig_to_new_prefix={
             # HF root is 'model.'
             "model.": "token2wav_model.",
-        })
+        }
+    )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -1508,6 +1734,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         # Provide placeholder to align with vLLM runner expectations
         def _empty_intermediate_tensors():
             return None
+
         self.make_empty_intermediate_tensors = _empty_intermediate_tensors
 
     def get_language_model(self) -> torch.nn.Module:
@@ -1555,7 +1782,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     ) -> Optional[SamplerOutput]:
         return None
 
-    def load_weights_without_buffers(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    def load_weights_without_buffers(
+        self, weights: Iterable[Tuple[str, torch.Tensor]]
+    ) -> Set[str]:
         loader = _Vllm_AutoWeightsLoader(self)
         loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
         # Log load summary
@@ -1567,7 +1796,11 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             device = next(self.parameters()).device
             self.logger.info(
                 "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
-                self.__class__.__name__, True, total_bytes / (1024**2), str(device))
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
         except Exception:
             pass
         return loaded
@@ -1575,29 +1808,27 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
     def find_all_registers(self):
         """
         Find all registered buffers in a PyTorch model.
-        
+
         Args:
         Returns:
             dict: Dictionary with buffer names as keys and their properties as values
         """
         registers = {}
-        
+
         # Get all named buffers
         for name, buf in self.named_buffers():
             if name in self.state_dict():
-                registers[name] = {
-                    'name': name,
-                    'buffer': buf
-                }
+                registers[name] = {"name": name, "buffer": buf}
         return registers
 
-        
-    #remove buffers from the weights and reload them after loading weights
-    def remove_buffers_from_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict):
+    # remove buffers from the weights and reload them after loading weights
+    def remove_buffers_from_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict
+    ):
         weights_to_load = []
         for key, value in weights:
             if key in buffers:
-                buffers[key]['buffer'] = value
+                buffers[key]["buffer"] = value
                 continue
             weights_to_load.append((key, value))
         return weights_to_load
@@ -1609,16 +1840,18 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         loaded_buffers = set()
         for name, buf_val in self.named_buffers():
             if name in buffers:
-                buf_val.copy_(buffers[name]['buffer'])
+                buf_val.copy_(buffers[name]["buffer"])
                 loaded_buffers.add(name)
         return loaded_buffers
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str) -> Set[str]:
+    def load_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str
+    ) -> Set[str]:
         buffers = self.find_all_registers()
         weights_to_load = self.remove_buffers_from_weights(weights, buffers)
         loaded = self.load_weights_without_buffers(weights_to_load)
         loaded_buffers = self.reload_buffers_to_model(buffers)
-        #merge loaded and loaded_buffers
+        # merge loaded and loaded_buffers
         loaded.update(loaded_buffers)
         self.spk_dict = torch.load(spk_dict_path)
         return loaded
@@ -1641,7 +1874,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
         )
 
     @torch.inference_mode()
-    def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> Optional[torch.Tensor]:
+    def process_chunk_bigvgan_batch(
+        self, mel_batch: torch.Tensor
+    ) -> Optional[torch.Tensor]:
         # BigVGAN is not part of this wrapper; return None for parity.
         return None
 
@@ -1684,7 +1919,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
             y_all=y_all,
             i=i,
             steps=steps,
-            prev_generated=prev_generated if isinstance(prev_generated, torch.Tensor) else None,
+            prev_generated=(
+                prev_generated if isinstance(prev_generated, torch.Tensor) else None
+            ),
             finished=finished,
-        )        
-        return _mel if _mel is not None else prev_generated, out
+        )
+        return _mel if _mel is not None else prev_generated, out

vllm_omni/model_executor/models/qwen2_old.py

@@ -3,7 +3,6 @@ from typing import Iterable, Optional, Set, Tuple, Union
 import torch
 from torch import nn
 from transformers import Qwen2Config
-
 from vllm.attention import Attention, AttentionType
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
@@ -11,27 +10,37 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
-                                               QKVParallelLinear,
-                                               RowParallelLinear)
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader, maybe_remap_kv_scale_name)
-from vllm.v1.pool.metadata import PoolingMetadata
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors, PoolerOutput
-
-from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
-                    is_pp_missing_parameter,
-                    make_empty_intermediate_tensors_factory, make_layers,
-                    maybe_prefix)
+from vllm.v1.pool.metadata import PoolingMetadata
 
 logger = init_logger(__name__)
 
@@ -62,8 +71,10 @@ class Qwen2MLP(nn.Module):
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
-            raise ValueError(f"Unsupported activation: {hidden_act}. "
-                             "Only silu is supported for now.")
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
@@ -75,18 +86,20 @@ class Qwen2MLP(nn.Module):
 
 class Qwen2Attention(nn.Module):
 
-    def __init__(self,
-                 hidden_size: int,
-                 num_heads: int,
-                 num_kv_heads: int,
-                 max_position: int = 4096 * 32,
-                 rope_theta: float = 10000,
-                 head_dim: Optional[int] = None,
-                 cache_config: Optional[CacheConfig] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 rope_scaling: Optional[Tuple] = None,
-                 prefix: str = "",
-                 attn_type: str = AttentionType.DECODER) -> None:
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        max_position: int = 4096 * 32,
+        rope_theta: float = 10000,
+        head_dim: Optional[int] = None,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        rope_scaling: Optional[Tuple] = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
         tp_size = get_tensor_model_parallel_world_size()
@@ -133,14 +146,16 @@ class Qwen2Attention(nn.Module):
             base=self.rope_theta,
             rope_scaling=rope_scaling,
         )
-        self.attn = Attention(self.num_heads,
-                              self.head_dim,
-                              self.scaling,
-                              num_kv_heads=self.num_kv_heads,
-                              cache_config=cache_config,
-                              quant_config=quant_config,
-                              prefix=f"{prefix}.attn",
-                              attn_type=attn_type)
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+        )
 
     def forward(
         self,
@@ -199,10 +214,10 @@ class Qwen2DecoderLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.mlp",
         )
-        self.input_layernorm = RMSNorm(config.hidden_size,
-                                       eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size,
-                                                eps=config.rms_norm_eps)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
 
     def forward(
         self,
@@ -215,16 +230,14 @@ class Qwen2DecoderLayer(nn.Module):
             residual = hidden_states
             hidden_states = self.input_layernorm(hidden_states)
         else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
-            hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 
@@ -237,14 +250,17 @@ class Qwen2DecoderLayer(nn.Module):
         "positions": -1,
         "intermediate_tensors": 0,
         "inputs_embeds": 0,
-    })
+    }
+)
 class Qwen2Model(nn.Module):
 
-    def __init__(self,
-                 *,
-                 vllm_config: VllmConfig,
-                 prefix: str = "",
-                 decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer,
+    ):
         super().__init__()
 
         config = vllm_config.model_config.hf_config
@@ -252,23 +268,27 @@ class Qwen2Model(nn.Module):
         quant_config = vllm_config.quant_config
 
         # TODO (@robertgshaw2): see if this can be moved out
-        if (cache_config.sliding_window is not None
-                and hasattr(config, "max_window_layers")):
-            raise ValueError("Sliding window for some but all layers is not "
-                             "supported. This model uses sliding window "
-                             "but `max_window_layers` = {} is less than "
-                             "`num_hidden_layers` = {}. Please open an issue "
-                             "to discuss this feature.".format(
-                                 config.max_window_layers,
-                                 config.num_hidden_layers,
-                             ))
+        if cache_config.sliding_window is not None and hasattr(
+            config, "max_window_layers"
+        ):
+            raise ValueError(
+                "Sliding window for some but all layers is not "
+                "supported. This model uses sliding window "
+                "but `max_window_layers` = {} is less than "
+                "`num_hidden_layers` = {}. Please open an issue "
+                "to discuss this feature.".format(
+                    config.max_window_layers,
+                    config.num_hidden_layers,
+                )
+            )
 
         self.config = config
         self.quant_config = quant_config
         self.vocab_size = config.vocab_size
 
-        if get_pp_group().is_first_rank or (config.tie_word_embeddings
-                                            and get_pp_group().is_last_rank):
+        if get_pp_group().is_first_rank or (
+            config.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 getattr(config, "embedding_size", config.hidden_size),
@@ -282,16 +302,18 @@ class Qwen2Model(nn.Module):
         decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: decoder_layer_type(config=config,
-                                              cache_config=cache_config,
-                                              quant_config=quant_config,
-                                              prefix=prefix),
+            lambda prefix: decoder_layer_type(
+                config=config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            ),
             prefix=f"{prefix}.layers",
         )
 
-        self.make_empty_intermediate_tensors = (
-            make_empty_intermediate_tensors_factory(
-                ["hidden_states", "residual"], config.hidden_size))
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
         if get_pp_group().is_last_rank:
             self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
@@ -317,22 +339,20 @@ class Qwen2Model(nn.Module):
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
-        for layer in self.layers[self.start_layer:self.end_layer]:
+        for layer in self.layers[self.start_layer : self.end_layer]:
             hidden_states, residual = layer(
                 positions,
                 hidden_states,
                 residual,
             )
         if not get_pp_group().is_last_rank:
-            return IntermediateTensors({
-                "hidden_states": hidden_states,
-                "residual": residual
-            })
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -346,18 +366,19 @@ class Qwen2Model(nn.Module):
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if (self.quant_config is not None and
-                (scale_name := self.quant_config.get_cache_scale(name))):
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(name)
+            ):
                 # Loading kv cache quantization scales
                 param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
-                                 loaded_weight[0])
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
                 weight_loader(param, loaded_weight)
                 loaded_params.add(scale_name)
                 continue
-            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)
@@ -381,8 +402,7 @@ class Qwen2Model(nn.Module):
                 if is_pp_missing_parameter(name, self):
                     continue
                 param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
             loaded_params.add(name)
         return loaded_params
@@ -411,18 +431,20 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config,
-                                prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         if get_pp_group().is_last_rank:
             if config.tie_word_embeddings:
                 self.lm_head = self.model.embed_tokens
             else:
-                self.lm_head = ParallelLMHead(config.vocab_size,
-                                              config.hidden_size,
-                                              quant_config=quant_config,
-                                              prefix=maybe_prefix(
-                                                  prefix, "lm_head"))
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
         else:
             self.lm_head = PPMissingLayer()
 
@@ -430,7 +452,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         self.sampler = get_sampler()
 
         self.make_empty_intermediate_tensors = (
-            self.model.make_empty_intermediate_tensors)
+            self.model.make_empty_intermediate_tensors
+        )
 
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
@@ -442,8 +465,9 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
-        hidden_states = self.model(input_ids, positions, intermediate_tensors,
-                                   inputs_embeds)
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
         return hidden_states
 
     def compute_logits(
@@ -451,8 +475,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         hidden_states: torch.Tensor,
         sampling_metadata: SamplingMetadata,
     ) -> Optional[torch.Tensor]:
-        logits = self.logits_processor(self.lm_head, hidden_states,
-                                       sampling_metadata)
+        logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
         return logits
 
     def sample(
@@ -463,12 +486,10 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
         next_tokens = self.sampler(logits, sampling_metadata)
         return next_tokens
 
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         loader = AutoWeightsLoader(
             self,
-            skip_prefixes=(["lm_head."]
-                           if self.config.tie_word_embeddings else None),
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
         return loader.load_weights(weights)
 
@@ -499,8 +520,9 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
         self.lora_config = lora_config
 
         self.quant_config = quant_config
-        self.model = Qwen2Model(vllm_config=vllm_config,
-                                prefix=maybe_prefix(prefix, "model"))
+        self.model = Qwen2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
 
         # TODO: Replace this model class with as_embedding_model(
         # Qwen2ForCausalLM) after changing the default pooling method
@@ -508,14 +530,13 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
             logger.warning(
                 "This embedding model will default to last-token pooling in "
                 "an upcoming version. To avoid breaking changes, you should "
-                "pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`"
-                " explicitly.")
+                'pass `--override-pooler-config \'{"pooling_type": "MEAN"}\'`'
+                " explicitly."
+            )
 
         self._pooler = Pooler.from_config_with_defaults(
-            pooler_config,
-            pooling_type=PoolingType.MEAN,
-            normalize=True,
-            softmax=False)
+            pooler_config, pooling_type=PoolingType.MEAN, normalize=True, softmax=False
+        )
 
     def forward(
         self,
@@ -534,6 +555,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = ((name, data) for name, data in weights
-                   if not name.startswith("lm_head."))
-        self.model.load_weights(weights)
+        weights = (
+            (name, data) for name, data in weights if not name.startswith("lm_head.")
+        )
+        self.model.load_weights(weights)

vllm_omni/model_executor/models/registry.py

@@ -0,0 +1,54 @@
+from vllm.model_executor.models.registry import (
+    _VLLM_MODELS,
+    _LazyRegisteredModel,
+    _ModelRegistry,
+)
+
+_OMNI_MODELS = {
+    "Qwen2_5OmniForConditionalGeneration": (
+        "qwen2_5_omni",
+        "Qwen2_5OmniForConditionalGeneration",
+    ),
+    "Qwen2_5OmniThinkerModel": (
+        "qwen2_5_omni_thinker",
+        "Qwen2_5OmniThinkerForConditionalGeneration",
+    ),  # noqa: E501
+    "Qwen2_5OmniTalkerModel": (
+        "qwen2_5_omni_talker",
+        "Qwen2_5OmniTalkerForConditionalGeneration",
+    ),  # noqa: E501
+    "Qwen2_5OmniToken2WavModel": (
+        "qwen2_5_omni_token2wav",
+        "Qwen2_5OmniToken2WavForConditionalGenerationVLLM",
+    ),
+    "Qwen2_5OmniToken2WavDiTModel": (
+        "qwen2_5_omni_token2wav",
+        "Qwen2_5OmniToken2WavModel",
+    ),
+    "Qwen2ForCausalLM_old": ("qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
+}
+
+_VLLM_OMNI_MODELS = {
+    **_VLLM_MODELS,
+    **_OMNI_MODELS,
+}
+
+
+OmniModelRegistry = _ModelRegistry(
+    {
+        **{
+            model_arch: _LazyRegisteredModel(
+                module_name=f"vllm.model_executor.models.{mod_relname}",
+                class_name=cls_name,
+            )
+            for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()
+        },
+        **{
+            model_arch: _LazyRegisteredModel(
+                module_name=f"vllm_omni.model_executor.models.{mod_relname}",
+                class_name=cls_name,
+            )
+            for model_arch, (mod_relname, cls_name) in _OMNI_MODELS.items()
+        },
+    }
+)

vllm_omni/model_executor/models/utils.py

@@ -5,4 +5,4 @@ def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
     """
     Add a prefix to the names of the loaded weights.
     """
-    return {maybe_prefix(prefix, name) for name in weights}
+    return {maybe_prefix(prefix, name) for name in weights}

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -37,4 +37,4 @@ stage_args:
       engine_output_type: audio
     engine_input_source: [1]
     final_output: true
-    final_output_type: audio
+    final_output_type: audio

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -1,40 +1,56 @@
 from typing import Union
-import torch
 
+import torch
 from vllm.inputs import TextPrompt
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
-def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
+
+def thinker2talker(
+    stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
+):
     if not engine_input_source:
-      raise ValueError("engine_input_source cannot be empty")
+        raise ValueError("engine_input_source cannot be empty")
     source_stage_id = engine_input_source[0]
     if source_stage_id >= len(stage_list):
-      raise IndexError(f"Invalid stage_id: {source_stage_id}")
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
     if stage_list[source_stage_id].engine_outputs is None:
         raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
     thinker_outputs = stage_list[source_stage_id].engine_outputs
     talker_inputs = []
-    multi_modal_data = {thinker_output.request_id: 
-            prompt.get('multi_modal_data', None) for thinker_output, prompt in zip(thinker_outputs, prompt)}
-    
+    multi_modal_data = {
+        thinker_output.request_id: prompt.get("multi_modal_data", None)
+        for thinker_output, prompt in zip(thinker_outputs, prompt)
+    }
+
     for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
         prompt_token_ids = thinker_output.prompt_token_ids
         thinker_output_ids = output.token_ids
         prompt_token_ids_len = len(prompt_token_ids)
-        thinker_hidden_states = output.multimodal_output["latent"].clone().detach().cuda()
+        thinker_hidden_states = (
+            output.multimodal_output["latent"].clone().detach().cuda()
+        )
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[0] * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token
+                prompt_token_ids=[0]
+                * (len(prompt_token_ids) + 2),  # add 2 for codec pad and start token
                 additional_information={
-                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
-                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
+                    "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(
+                        torch.float32
+                    ),
+                    "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(
+                        torch.float32
+                    ),
                     "prompt_token_ids": prompt_token_ids,
                     "thinker_output_token_ids": thinker_output_ids,
                 },
-                multi_modal_data=multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None,
+                multi_modal_data=(
+                    multi_modal_data[thinker_output.request_id]
+                    if multi_modal_data is not None
+                    else None
+                ),
                 mm_processor_kwargs=None,
             )
         )
-    return talker_inputs
+    return talker_inputs

vllm_omni/outputs.py

@@ -1,15 +1,17 @@
-from vllm.v1.outputs import ModelRunnerOutput
-from vllm.outputs import RequestOutput
-from typing import Optional
 from dataclasses import dataclass
+from typing import Optional
+
 import torch
+from vllm.outputs import RequestOutput
+from vllm.v1.outputs import ModelRunnerOutput
 
 
 class OmniModelRunnerOutput(ModelRunnerOutput):
     multimodal_outputs: Optional[dict[str, torch.Tensor]] = None
 
+
 @dataclass
 class OmniRequestOutput(RequestOutput):
     stage_id: int
     final_output_type: str
-    request_output: RequestOutput
+    request_output: RequestOutput

vllm_omni/patch.py

@@ -1,21 +1,29 @@
 import sys
 
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
-from vllm.v1.request import Request as _OriginalRequest
+from vllm.model_executor.layers.rotary_embedding import (
+    MRotaryEmbedding as _OriginalMRotaryEmbedding,
+)
 from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
+from vllm.v1.request import Request as _OriginalRequest
 
+from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 from vllm_omni.request import OmniRequest
-from vllm_omni.engine import OmniEngineCoreRequest
 
 for module_name, module in sys.modules.items():
-    if hasattr(module, 'TokensPrompt') and module.TokensPrompt == _OriginalTokensPrompt:
+    if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
         module.TokensPrompt = OmniTokensPrompt
-    if hasattr(module, 'MRotaryEmbedding') and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
+    if (
+        hasattr(module, "MRotaryEmbedding")
+        and module.MRotaryEmbedding == _OriginalMRotaryEmbedding
+    ):
         module.MRotaryEmbedding = MRotaryEmbedding
-    if hasattr(module, 'Request') and module.Request == _OriginalRequest:
+    if hasattr(module, "Request") and module.Request == _OriginalRequest:
         module.Request = OmniRequest
-    if hasattr(module, 'EngineCoreRequest') and module.EngineCoreRequest == _OriginalEngineCoreRequest:
-        module.EngineCoreRequest = OmniEngineCoreRequest
+    if (
+        hasattr(module, "EngineCoreRequest")
+        and module.EngineCoreRequest == _OriginalEngineCoreRequest
+    ):
+        module.EngineCoreRequest = OmniEngineCoreRequest

vllm_omni/request.py

@@ -1,29 +1,40 @@
-from vllm.v1.request import Request
-from vllm.v1.structured_output.request import StructuredOutputRequest
+from typing import Optional
+
 from vllm.multimodal.inputs import MultiModalKwargs
 from vllm.utils import is_list_of
-from typing import Optional
-from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload, OmniEngineCoreRequest
+from vllm.v1.request import Request
+from vllm.v1.structured_output.request import StructuredOutputRequest
+
+from vllm_omni.engine import (
+    AdditionalInformationPayload,
+    OmniEngineCoreRequest,
+    PromptEmbedsPayload,
+)
 
 
 class OmniRequest(Request):
     def __init__(
-                 self, 
-                 prompt_embeds: Optional[PromptEmbedsPayload] = None,
-                 additional_information: Optional[AdditionalInformationPayload] = None,
-                 *args, **kwargs):
+        self,
+        prompt_embeds: Optional[PromptEmbedsPayload] = None,
+        additional_information: Optional[AdditionalInformationPayload] = None,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         # Serialized prompt embeddings payload (optional)
         self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds
         # Serialized additional information payload (optional)
-        self.additional_information: Optional[AdditionalInformationPayload] = additional_information
-    
+        self.additional_information: Optional[AdditionalInformationPayload] = (
+            additional_information
+        )
+
     @classmethod
     def from_engine_core_request(cls, request: OmniEngineCoreRequest) -> "Request":
         if request.mm_inputs is not None:
             assert isinstance(request.mm_inputs, list)
-            assert is_list_of(request.mm_inputs, MultiModalKwargs), (
-                "mm_inputs was not updated in EngineCore.add_request")
+            assert is_list_of(
+                request.mm_inputs, MultiModalKwargs
+            ), "mm_inputs was not updated in EngineCore.add_request"
 
         return cls(
             request_id=request.request_id,
@@ -37,11 +48,13 @@ class OmniRequest(Request):
             eos_token_id=request.eos_token_id,
             arrival_time=request.arrival_time,
             lora_request=request.lora_request,
-            structured_output_request=StructuredOutputRequest(
-                sampling_params=request.sampling_params) \
-                    if request.sampling_params else None,
+            structured_output_request=(
+                StructuredOutputRequest(sampling_params=request.sampling_params)
+                if request.sampling_params
+                else None
+            ),
             cache_salt=request.cache_salt,
             priority=request.priority,
             prompt_embeds=request.prompt_embeds,
             additional_information=request.additional_information,
-        )
+        )

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,12 +6,15 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Any, List
+from typing import Any, List, Optional, Union
+
 import numpy as np
-
 import torch
-
 from vllm import envs
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.worker.gpu_model_runner import (
     EMPTY_MODEL_RUNNER_OUTPUT,
     IntermediateTensors,
@@ -20,12 +23,8 @@ from vllm.v1.worker.gpu_model_runner import (
     has_kv_transfer_group,
     set_forward_context,
 )
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.multimodal.inputs import MultiModalKwargs
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
-from vllm.v1.spec_decode.eagle import EagleProposer
 
-from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload
+from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
@@ -42,9 +41,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.input_ids = torch.zeros(self.max_num_tokens,
-                                     dtype=torch.int64,
-                                     device=self.device)
+        self.input_ids = torch.zeros(
+            self.max_num_tokens, dtype=torch.int64, device=self.device
+        )
 
     @torch.inference_mode()
     def execute_model(
@@ -55,12 +54,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
         # Update internal state with the new schedule
         self._update_states(scheduler_output)
 
-        # Decode per-request prompt_embeds / additional_hidden_states payloads (if present) into CPU tensors
+        # Decode per-request prompt_embeds / additional_hidden_states payloads
+        # (if present) into CPU tensors
         try:
             new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
             if new_reqs:
                 for nr in new_reqs:
-                    req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+                    req_id = getattr(nr, "req_id", None) or getattr(
+                        nr, "request_id", None
+                    )
                     if req_id is None:
                         continue
                     # prompt_embeds
@@ -87,14 +89,20 @@ class GPUARModelRunner(OmniGPUModelRunner):
                         elif isinstance(payload_info, AdditionalInformationPayload):
                             for k, entry in payload_info.entries.items():
                                 if entry.tensor_data is not None:
-                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+                                    dt = np.dtype(
+                                        getattr(entry, "tensor_dtype", "float32")
+                                    )
                                     arr = np.frombuffer(entry.tensor_data, dtype=dt)
                                     arr = arr.reshape(entry.tensor_shape)
                                     info_dict[k] = torch.from_numpy(arr)
                                 else:
                                     info_dict[k] = entry.list_data
                         if info_dict and req_id in self.requests:
-                            setattr(self.requests[req_id], "additional_information_cpu", info_dict)
+                            setattr(
+                                self.requests[req_id],
+                                "additional_information_cpu",
+                                info_dict,
+                            )
         except Exception:
             pass
 
@@ -102,26 +110,33 @@ class GPUARModelRunner(OmniGPUModelRunner):
         if not scheduler_output.total_num_scheduled_tokens:
             if not has_kv_transfer_group():
                 return EMPTY_MODEL_RUNNER_OUTPUT
-            return self.kv_connector_no_forward(scheduler_output,
-                                                self.vllm_config)
+            return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
 
         # Prepare decoder inputs and attention metadata
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = self._prepare_inputs(
-             scheduler_output)
+        (
+            attn_metadata,
+            attention_cuda_graphs,
+            logits_indices,
+            spec_decode_metadata,
+            num_scheduled_tokens_np,
+            spec_decode_common_attn_metadata,
+        ) = self._prepare_inputs(scheduler_output)
 
         # Determine number of input tokens for this iteration
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
+        if (
+            self.use_cuda_graph
+            and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]
+        ):
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)
         else:
             tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if (self.compilation_config.pass_config.enable_sequence_parallelism
-                    and tp_size > 1):
+            if (
+                self.compilation_config.pass_config.enable_sequence_parallelism
+                and tp_size > 1
+            ):
                 from vllm.utils import round_up  # lazy local import
+
                 num_input_tokens = round_up(num_scheduled_tokens, tp_size)
             else:
                 num_input_tokens = num_scheduled_tokens
@@ -137,29 +152,32 @@ class GPUARModelRunner(OmniGPUModelRunner):
         else:
             mm_embeds = []
 
-        # Always assemble inputs_embeds on first PP rank; overlay per-request prompt_embeds and collect additional_hidden_states for prefill only
+        # Always assemble inputs_embeds on first PP rank; overlay per-request
+        # prompt_embeds and collect additional_hidden_states for prefill only
         if get_pp_group().is_first_rank:
             inputs_embeds_scheduled = self.model.get_input_embeddings(
                 input_ids=self.input_ids[:num_scheduled_tokens],
-                multimodal_embeddings=(mm_embeds or None)
-                if self.is_multimodal_model else None,
+                multimodal_embeddings=(
+                    (mm_embeds or None) if self.is_multimodal_model else None
+                ),
             )
 
             # Copy into persistent buffer to enable CUDA Graph capture
-            self.inputs_embeds[:num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
+            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
 
             # Reset per-step additional information collector
             if hasattr(self, "_forward_additional_information"):
                 self._forward_additional_information = None
 
-            # Overlay custom prompt_embeds per request for the prompt portion; collect additional_information (tensor/list) for prefill portion only
+            # Overlay custom prompt_embeds per request for the prompt portion;
+            # collect additional_information (tensor/list) for prefill portion only
             for req_index, req_id in enumerate(self.input_batch.req_ids):
                 req_state = self.requests[req_id]
                 pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
                 addi_cpu = getattr(req_state, "additional_information_cpu", None)
                 num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index])
+                    self.input_batch.num_computed_tokens_cpu[req_index]
+                )
                 prompt_len = len(req_state.prompt_token_ids)
                 prompt_remaining = max(0, prompt_len - num_computed_tokens)
                 sched_tokens = int(num_scheduled_tokens_np[req_index])
@@ -167,39 +185,48 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 if overlay_len <= 0:
                     continue
                 if pe_cpu is not None:
-                    src = pe_cpu[num_computed_tokens:
-                                 num_computed_tokens + overlay_len].to(
-                                      dtype=self.dtype,
-                                      device=self.device,
-                                      non_blocking=True)
+                    src = pe_cpu[
+                        num_computed_tokens : num_computed_tokens + overlay_len
+                    ].to(dtype=self.dtype, device=self.device, non_blocking=True)
                     start_offset = int(self.query_start_loc_cpu[req_index])
-                    self.inputs_embeds[start_offset:start_offset + overlay_len] \
-                        .copy_(src)
+                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(
+                        src
+                    )
                 # For additional_information: handle arbitrary keys
                 if addi_cpu is not None and isinstance(addi_cpu, dict):
                     # Lazy init collector dict
-                    if not hasattr(self, "_forward_additional_information") or \
-                       self._forward_additional_information is None:
+                    if (
+                        not hasattr(self, "_forward_additional_information")
+                        or self._forward_additional_information is None
+                    ):
                         self._forward_additional_information = {}
-                    # Process tensors (slice by scheduled prompt range) and lists (append per-request)
+                    # Process tensors (slice by scheduled prompt range) and
+                    # lists (append per-request)
                     for k, v in addi_cpu.items():
                         if isinstance(v, torch.Tensor):
                             # Slice along token dimension for prefill part
                             try:
-                                seg = v[num_computed_tokens:
-                                        num_computed_tokens + overlay_len].to(
-                                            dtype=self.dtype,
-                                            device=self.device,
-                                            non_blocking=True)
+                                seg = v[
+                                    num_computed_tokens : num_computed_tokens
+                                    + overlay_len
+                                ].to(
+                                    dtype=self.dtype,
+                                    device=self.device,
+                                    non_blocking=True,
+                                )
                             except Exception:
                                 # Fallback: move whole tensor if slicing fails
-                                seg = v.to(dtype=self.dtype,
-                                           device=self.device,
-                                           non_blocking=True)
+                                seg = v.to(
+                                    dtype=self.dtype,
+                                    device=self.device,
+                                    non_blocking=True,
+                                )
                             prev_val = self._forward_additional_information.get(k)
                             self._forward_additional_information[k] = (
                                 torch.cat([prev_val, seg], dim=0)
-                                if isinstance(prev_val, torch.Tensor) else seg.clone())
+                                if isinstance(prev_val, torch.Tensor)
+                                else seg.clone()
+                            )
                         elif isinstance(v, list):
                             prev_val = self._forward_additional_information.get(k)
                             if prev_val is None:
@@ -212,8 +239,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
             input_ids = self.input_ids[:num_input_tokens]  # preserved for APIs
             inputs_embeds = self.inputs_embeds[:num_input_tokens]
-            model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
-                               if self.is_multimodal_model else {})
+            model_mm_kwargs = (
+                self._extract_mm_kwargs(scheduler_output)
+                if self.is_multimodal_model
+                else {}
+            )
         else:
             input_ids = self.input_ids[:num_input_tokens]
             inputs_embeds = None
@@ -230,27 +260,31 @@ class GPUARModelRunner(OmniGPUModelRunner):
             intermediate_tensors = None
         else:
             intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
+                num_input_tokens, intermediate_tensors, True
+            )
 
         # Some attention backends only support CUDA Graphs in pure decode.
         skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
 
         # Forward pass
         with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=num_input_tokens,
+            num_tokens_across_dp=num_tokens_across_dp,
+            skip_cuda_graphs=skip_cuda_graphs,
+        ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output:
 
             model_kwargs_extra = {}
             # Only pass additional_information for the prefill part
-            if hasattr(self, "_forward_additional_information") and \
-               self._forward_additional_information is not None and \
-               isinstance(self._forward_additional_information, dict):
-                model_kwargs_extra["additional_information"] = self._forward_additional_information
+            if (
+                hasattr(self, "_forward_additional_information")
+                and self._forward_additional_information is not None
+                and isinstance(self._forward_additional_information, dict)
+            ):
+                model_kwargs_extra["additional_information"] = (
+                    self._forward_additional_information
+                )
             model_output = self.model(
                 input_ids=input_ids,
                 positions=positions,
@@ -272,8 +306,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
             hidden_states = model_output
             _aux_hidden_states = None
 
-        text_hidden_states, multimodal_outputs = (
-            self.extract_multimodal_outputs(hidden_states))
+        text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            hidden_states
+        )
 
         # Mid PP stages return intermediate tensors unmodified
         if not get_pp_group().is_last_rank:
@@ -282,30 +317,41 @@ class GPUARModelRunner(OmniGPUModelRunner):
             return text_hidden_states
 
         # Broadcast PP output for external_launcher (torchrun)
-        broadcast_pp_output = \
-            self.parallel_config.distributed_executor_backend \
-            == "external_launcher" and len(get_pp_group().ranks) > 0
+        broadcast_pp_output = (
+            self.parallel_config.distributed_executor_backend == "external_launcher"
+            and len(get_pp_group().ranks) > 0
+        )
         if not get_pp_group().is_last_rank:
             assert isinstance(text_hidden_states, IntermediateTensors)
             if not broadcast_pp_output:
                 text_hidden_states.kv_connector_output = kv_connector_output
                 return text_hidden_states
-            get_pp_group().send_tensor_dict(text_hidden_states.tensors,
-                                            all_gather_group=get_tp_group())
+            get_pp_group().send_tensor_dict(
+                text_hidden_states.tensors, all_gather_group=get_tp_group()
+            )
             logits = None
         else:
             if self.input_batch.pooling_params:
-                return self._pool(text_hidden_states, num_scheduled_tokens,
-                                  num_scheduled_tokens_np, kv_connector_output)
+                return self._pool(
+                    text_hidden_states,
+                    num_scheduled_tokens,
+                    num_scheduled_tokens_np,
+                    kv_connector_output,
+                )
 
             sample_hidden_states = text_hidden_states[logits_indices]
             logits = self.model.compute_logits(sample_hidden_states, None)
         if broadcast_pp_output:
-            model_output_broadcast_data = {
-                "logits": logits.contiguous(),
-            } if logits is not None else {}
+            model_output_broadcast_data = (
+                {
+                    "logits": logits.contiguous(),
+                }
+                if logits is not None
+                else {}
+            )
             model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
-                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
+                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
+            )
             assert model_output_broadcast_data is not None
             logits = model_output_broadcast_data["logits"]
 
@@ -346,8 +392,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
         discard_sampled_tokens_req_indices = []
         for i, req_id in enumerate(self.input_batch.req_ids):
             req_state = self.requests[req_id]
-            seq_len = (req_state.num_computed_tokens +
-                       scheduler_output.num_scheduled_tokens[req_id])
+            seq_len = (
+                req_state.num_computed_tokens
+                + scheduler_output.num_scheduled_tokens[req_id]
+            )
             if seq_len < req_state.num_tokens:
                 generator = self.input_batch.generators.get(i)
                 if generator is not None:
@@ -356,8 +404,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # Move CPU sync parts
         logprobs_tensors = sampler_output.logprobs_tensors
-        logprobs_lists = logprobs_tensors.tolists() \
-            if logprobs_tensors is not None else None
+        logprobs_lists = (
+            logprobs_tensors.tolists() if logprobs_tensors is not None else None
+        )
 
         # Prompt logprobs if needed
         prompt_logprobs_dict = self._get_prompt_logprobs_dict(
@@ -412,10 +461,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
         pooler_output: list[Optional[torch.Tensor]] = []
         prev_logits_index = 0
         for logits_index in logits_indices:
-            pooler_output.append(text_hidden_states_cpu[prev_logits_index:logits_index+1])
+            pooler_output.append(
+                text_hidden_states_cpu[prev_logits_index : logits_index + 1]
+            )
             prev_logits_index = logits_index + 1
 
-
         self.eplb_step()
 
         return OmniModelRunnerOutput(
@@ -425,14 +475,23 @@ class GPUARModelRunner(OmniGPUModelRunner):
             spec_token_ids=spec_token_ids,
             logprobs=logprobs_lists,
             prompt_logprobs_dict=prompt_logprobs_dict,
-            pooler_output=pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None,
+            pooler_output=(
+                pooler_output
+                if self.vllm_config.model_config.engine_output_type != "text"
+                else None
+            ),
             kv_connector_output=kv_connector_output,
             num_nans_in_logits=num_nans_in_logits,
         )
-    
+
     @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]) -> dict:
-        if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
+    def extract_multimodal_outputs(
+        self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]
+    ) -> dict:
+        if (
+            hasattr(self.model, "have_multimodal_outputs")
+            and self.model.have_multimodal_outputs
+        ):
             text_hidden_states = hidden_states.text_hidden_states
             multimodal_outputs = hidden_states.multimodal_outputs
 
@@ -445,7 +504,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
         else:
             raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
         return text_hidden_states, multimodal_outputs
-    
+
     @torch.inference_mode()
     def _dummy_run(
         self,
@@ -470,8 +529,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
         num_scheduled_tokens_list[-1] += num_tokens % num_reqs
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
 
         attn_metadata: Optional[dict[str, Any]] = None
         if capture_attn_cudagraph:
@@ -480,36 +538,43 @@ class GPUARModelRunner(OmniGPUModelRunner):
             # Make sure max_model_len is used at the graph capture time.
             self.seq_lens_np[:num_reqs] = self.max_model_len
             self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
+            self.seq_lens[:num_reqs].copy_(
+                self.seq_lens_cpu[:num_reqs], non_blocking=True
+            )
 
             for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
+                self.kv_cache_config.kv_cache_groups
+            ):
                 common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
-                                                                 1],
+                    query_start_loc=self.query_start_loc[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
                     seq_lens=self.seq_lens[:num_reqs],
                     seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
                     num_reqs=num_reqs,
                     num_actual_tokens=num_tokens,
                     max_query_len=num_tokens,
                     block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
-                    slot_mapping=self.input_batch.
-                    block_table[kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
+                        kv_cache_group_id
+                    ].get_device_tensor()[:num_reqs],
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping[:num_tokens],
+                    causal=True,
+                )
 
                 for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder\
-                        .build_for_cudagraph_capture(common_attn_metadata)
+                    attn_metadata_i = (
+                        attn_group.metadata_builder.build_for_cudagraph_capture(
+                            common_attn_metadata
+                        )
+                    )
                     for layer_name in kv_cache_group_spec.layer_names:
                         attn_metadata[layer_name] = attn_metadata_i
 
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens):
             if self.is_multimodal_model:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds[:num_tokens]
@@ -532,16 +597,20 @@ class GPUARModelRunner(OmniGPUModelRunner):
                         self.model.make_empty_intermediate_tensors(
                             batch_size=self.max_num_tokens,
                             dtype=self.model_config.dtype,
-                            device=self.device))
+                            device=self.device,
+                        )
+                    )
 
                 intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
+                    num_tokens, None, False
+                )
 
             with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+            ):
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,
@@ -573,6 +642,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             self.eplb_step(is_dummy=True, is_profile=is_profile)
 
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+        hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            hidden_states
+        )
         return hidden_states, hidden_states[logit_indices]
-

vllm_omni/worker/gpu_ar_worker.py

@@ -1,13 +1,14 @@
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-import torch
-import os
 import gc
-from vllm.utils import GiB_bytes, MemorySnapshot
-from vllm.platforms import current_platform
-from vllm.worker.worker import _check_if_gpu_supports_dtype
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+import os
+
+import torch
 from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.worker.worker import _check_if_gpu_supports_dtype
 
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
@@ -34,10 +35,15 @@ class GPUARWorker(GPUWorker):
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (self.init_snapshot.total_memory *
-                                     self.cache_config.gpu_memory_utilization)
+            self.requested_memory = (
+                self.init_snapshot.total_memory
+                * self.cache_config.gpu_memory_utilization
+            )
             if self.init_snapshot.free_memory < self.requested_memory:
-                GiB = lambda b: round(b / GiB_bytes, 2)
+
+                def GiB(b):
+                    return round(b / GiB_bytes, 2)
+
                 raise ValueError(
                     f"Free memory on device "
                     f"({GiB(self.init_snapshot.free_memory)}/"
@@ -48,20 +54,23 @@ class GPUARWorker(GPUWorker):
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
-            raise RuntimeError(
-                f"Not support device type: {self.device_config.device}")
+            raise RuntimeError(f"Not support device type: {self.device_config.device}")
         # Initialize the distributed environment.
-        init_worker_distributed_environment(self.vllm_config, self.rank,
-                                            self.distributed_init_method,
-                                            self.local_rank,
-                                            current_platform.dist_backend)
+        init_worker_distributed_environment(
+            self.vllm_config,
+            self.rank,
+            self.distributed_init_method,
+            self.local_rank,
+            current_platform.dist_backend,
+        )
         # Set random seed.
         set_random_seed(self.model_config.seed)
 
         # Construct the model runner
         self.model_runner: GPUARModelRunner = GPUARModelRunner(
-            self.vllm_config, self.device)
+            self.vllm_config, self.device
+        )
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
-            report_usage_stats(self.vllm_config)
+            report_usage_stats(self.vllm_config)

vllm_omni/worker/gpu_diffusion_model_runner.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-from typing import Optional, Union, List
-import numpy as np
 import gc
 import logging
+from typing import List, Optional, Union
 
+import numpy as np
 import torch
-
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu_model_runner import (
     EMPTY_MODEL_RUNNER_OUTPUT,
     IntermediateTensors,
@@ -14,15 +16,10 @@ from vllm.v1.worker.gpu_model_runner import (
     set_forward_context,
 )
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
-from vllm.multimodal.inputs import MultiModalKwargs
-
-from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -45,12 +42,16 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         if not scheduler_output.total_num_scheduled_tokens:
 
             return EMPTY_MODEL_RUNNER_OUTPUT
-           
+
         # Prepare decoder inputs and attention metadata (for batch/order mapping)
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = self._prepare_inputs(
-             scheduler_output)
+        (
+            attn_metadata,
+            attention_cuda_graphs,
+            logits_indices,
+            spec_decode_metadata,
+            num_scheduled_tokens_np,
+            spec_decode_common_attn_metadata,
+        ) = self._prepare_inputs(scheduler_output)
 
         # Input token count for this iteration (not used by diffusion, but
         # retained to keep DP padding/ordering consistent)
@@ -68,11 +69,12 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         # Build inputs to mirror AR runner: input_ids/positions/embeds
         if self.is_multimodal_model and get_pp_group().is_first_rank:
             inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids[:scheduler_output.total_num_scheduled_tokens],
+                input_ids=self.input_ids[: scheduler_output.total_num_scheduled_tokens],
                 multimodal_embeddings=mm_embeds or None,
             )
-            self.inputs_embeds[:scheduler_output.total_num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
+            self.inputs_embeds[: scheduler_output.total_num_scheduled_tokens].copy_(
+                inputs_embeds_scheduled
+            )
             input_ids = self.input_ids[:num_input_tokens]
             inputs_embeds = self.inputs_embeds[:num_input_tokens]
             model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
@@ -92,18 +94,18 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             intermediate_tensors = None
         else:
             intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
+                num_input_tokens, intermediate_tensors, True
+            )
 
         # Set forward context mainly for resource management and kv connector
         skip_cuda_graphs = True  # diffusion path does not rely on cuda graphs here
         with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=num_input_tokens,
+            num_tokens_across_dp=num_tokens_across_dp,
+            skip_cuda_graphs=skip_cuda_graphs,
+        ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output:
 
             if not get_pp_group().is_last_rank:
                 # For non-last PP stages, pass through intermediate tensors.
@@ -119,8 +121,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                 multimodal_kwargs=model_mm_kwargs,
                 logits_indices=logits_indices,
             )
-        _, multimodal_outputs = (
-            self.extract_multimodal_outputs(outputs))
+        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
 
         # Ensure one tensor per request, map to CPU for output struct
         pooler_output: List[Optional[torch.Tensor]] = []
@@ -128,13 +129,19 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             # If model returned a single stacked tensor, split by requests
             assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+                pooler_output.append(
+                    multimodal_outputs[i].detach().to("cpu").contiguous()
+                )
         elif isinstance(multimodal_outputs, list):
             for out in multimodal_outputs:
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
         elif isinstance(multimodal_outputs, dict):
             for out in multimodal_outputs.values():
-                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+                pooler_output.append(
+                    out.detach().to("cpu").contiguous() if out is not None else None
+                )
         else:
             raise RuntimeError("Unsupported diffusion output type")
 
@@ -184,12 +191,12 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         # For Qwen 2.5 Omni's current implementation, we only support the forward method
         if hasattr(self.model, "forward"):
             return self.model.forward(**kwargs)
-        #TODO: add the diffuse method for other models
+        # TODO: add the diffuse method for other models
 
         raise RuntimeError(
             "The loaded model does not expose diffusion interfaces 'sample', "
-            "'forward', or 'diffuse'. Please implement one of them or adapt the runner.")
-
+            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+        )
 
     @torch.inference_mode()
     def _dummy_run(
@@ -214,8 +221,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
         num_scheduled_tokens_list[-1] += num_tokens % num_reqs
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
 
         attn_metadata: Optional[dict[str, dict]] = None
         if capture_attn_cudagraph:
@@ -224,35 +230,43 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             # Make sure max_model_len is used at the graph capture time.
             self.seq_lens_np[:num_reqs] = self.max_model_len
             self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
+            self.seq_lens[:num_reqs].copy_(
+                self.seq_lens_cpu[:num_reqs], non_blocking=True
+            )
 
             for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
+                self.kv_cache_config.kv_cache_groups
+            ):
                 common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
+                    query_start_loc=self.query_start_loc[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
                     seq_lens=self.seq_lens[:num_reqs],
                     seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
                     num_reqs=num_reqs,
                     num_actual_tokens=num_tokens,
                     max_query_len=num_tokens,
                     block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
+                        kv_cache_group_id
+                    ].get_device_tensor()[:num_reqs],
                     slot_mapping=self.input_batch.block_table[
-                        kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
+                        kv_cache_group_id
+                    ].slot_mapping[:num_tokens],
+                    causal=True,
+                )
 
                 for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder \
-                        .build_for_cudagraph_capture(common_attn_metadata)
+                    attn_metadata_i = (
+                        attn_group.metadata_builder.build_for_cudagraph_capture(
+                            common_attn_metadata
+                        )
+                    )
                     for layer_name in kv_cache_group_spec.layer_names:
                         attn_metadata[layer_name] = attn_metadata_i
 
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens):
             if self.is_multimodal_model:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds[:num_tokens]
@@ -275,17 +289,21 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                         self.model.make_empty_intermediate_tensors(
                             batch_size=self.max_num_tokens,
                             dtype=self.model_config.dtype,
-                            device=self.device))
+                            device=self.device,
+                        )
+                    )
 
                 intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
+                    num_tokens, None, False
+                )
 
             # Diffusion path: avoid CUDA graphs; we only use context for resource wiring
             with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+            ):
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,
@@ -311,9 +329,9 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             if not skip_eplb:
                 self.eplb_step(is_dummy=True, is_profile=is_profile)
 
-        logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        # logit_indices = np.cumsum(num_scheduled_tokens) - 1  # unused variable
         return text_hidden_states, None
-    
+
     @torch.inference_mode()
     def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
         logger.warning("Dummy sampler run is not implemented for diffusion model")
@@ -327,7 +345,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
             assert mm_budget is not None
 
             # TODO: handle encoder-decoder models once supported.
-            if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+            if (mm_budget.get_encoder_budget()) > 0:  # encoder_budget unused
                 (
                     dummy_modality,
                     max_tokens,
@@ -343,15 +361,15 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
                 )
 
                 dummy_encoder_outputs = self.model.get_multimodal_embeddings(
-                    **batched_dummy_mm_inputs)
+                    **batched_dummy_mm_inputs
+                )
 
                 sanity_check_mm_encoder_outputs(
                     dummy_encoder_outputs,
                     expected_num_items=max_mm_items_per_batch,
                 )
 
-                self.encoder_cache["tmp"] = dict(
-                    enumerate(dummy_encoder_outputs))
+                self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
 
         hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
         if get_pp_group().is_last_rank:

vllm_omni/worker/gpu_diffusion_worker.py

@@ -1,13 +1,14 @@
-from vllm.v1.worker.gpu_worker import Worker as GPUWorker
-import torch
-import os
 import gc
-from vllm.utils import GiB_bytes, MemorySnapshot
-from vllm.platforms import current_platform
-from vllm.worker.worker import _check_if_gpu_supports_dtype
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+import os
+
+import torch
 from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
 from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.worker.worker import _check_if_gpu_supports_dtype
 
 from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
@@ -34,10 +35,15 @@ class GPUDiffusionWorker(GPUWorker):
 
             # take current memory snapshot
             self.init_snapshot = MemorySnapshot()
-            self.requested_memory = (self.init_snapshot.total_memory *
-                                     self.cache_config.gpu_memory_utilization)
+            self.requested_memory = (
+                self.init_snapshot.total_memory
+                * self.cache_config.gpu_memory_utilization
+            )
             if self.init_snapshot.free_memory < self.requested_memory:
-                GiB = lambda b: round(b / GiB_bytes, 2)
+
+                def GiB(b):
+                    return round(b / GiB_bytes, 2)
+
                 raise ValueError(
                     f"Free memory on device "
                     f"({GiB(self.init_snapshot.free_memory)}/"
@@ -48,20 +54,23 @@ class GPUDiffusionWorker(GPUWorker):
                     f"utilization or reduce GPU memory used by other processes."
                 )
         else:
-            raise RuntimeError(
-                f"Not support device type: {self.device_config.device}")
+            raise RuntimeError(f"Not support device type: {self.device_config.device}")
         # Initialize the distributed environment.
-        init_worker_distributed_environment(self.vllm_config, self.rank,
-                                            self.distributed_init_method,
-                                            self.local_rank,
-                                            current_platform.dist_backend)
+        init_worker_distributed_environment(
+            self.vllm_config,
+            self.rank,
+            self.distributed_init_method,
+            self.local_rank,
+            current_platform.dist_backend,
+        )
         # Set random seed.
         set_random_seed(self.model_config.seed)
 
         # Construct the model runner
         self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
-            self.vllm_config, self.device)
+            self.vllm_config, self.device
+        )
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
-            report_usage_stats(self.vllm_config)
+            report_usage_stats(self.vllm_config)

vllm_omni/worker/gpu_model_runner.py

@@ -1,37 +1,35 @@
-
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 import numpy as np
 import torch
-
 import vllm.envs as envs
+from vllm.distributed.kv_transfer import has_kv_transfer_group
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.sampling_params import SamplingType
 from vllm.sequence import IntermediateTensors
-from vllm.utils import LazyLoader
+from vllm.utils import LazyLoader, round_up
 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.sampling_params import SamplingType
-from vllm.distributed.kv_transfer import has_kv_transfer_group
-from vllm.utils import round_up
 from vllm.v1.spec_decode.eagle import EagleProposer
-
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
-from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
+from vllm_omni.outputs import OmniModelRunnerOutput
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
 else:
     xgr = LazyLoader("xgr", globals(), "xgrammar")
     xgr_torch_compile = LazyLoader(
-        "xgr_torch_compile", globals(),
-        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile")
+        "xgr_torch_compile",
+        globals(),
+        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile",
+    )
 
 logger = init_logger(__name__)
 
@@ -90,16 +88,19 @@ class OmniGPUModelRunner(GPUModelRunner):
             sampling_params = new_req_data.sampling_params
             pooling_params = new_req_data.pooling_params
 
-            if sampling_params and \
-                sampling_params.sampling_type == SamplingType.RANDOM_SEED:
+            if (
+                sampling_params
+                and sampling_params.sampling_type == SamplingType.RANDOM_SEED
+            ):
                 generator = torch.Generator(device=self.device)
                 generator.manual_seed(sampling_params.seed)
             else:
                 generator = None
 
             if pooling_params:
-                assert (task := pooling_params.task) is not None, (
-                    "You did not set `task` in the API")
+                assert (
+                    task := pooling_params.task
+                ) is not None, "You did not set `task` in the API"
 
                 model = cast(VllmModelForPooling, self.model)
                 to_update = model.pooler.get_pooling_updates(task)
@@ -129,7 +130,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                     pe_cpu = torch.from_numpy(arr)
                     # Store temporarily on CPU; later moved to device in builder
                     setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
-                    # Also replace payload with Tensor for user visibility in scheduler_output
+                    # Also replace payload with Tensor for user visibility in
+                    # scheduler_output
                     try:
                         new_req_data.prompt_embeds = pe_cpu  # type: ignore[assignment]
                     except Exception:
@@ -145,17 +147,24 @@ class OmniGPUModelRunner(GPUModelRunner):
                         info_dict = payload_info
                     else:
                         from vllm.v1.engine import AdditionalInformationPayload
+
                         if isinstance(payload_info, AdditionalInformationPayload):
                             for k, entry in payload_info.entries.items():
                                 if entry.tensor_data is not None:
-                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+                                    dt = np.dtype(
+                                        getattr(entry, "tensor_dtype", "float32")
+                                    )
                                     arr = np.frombuffer(entry.tensor_data, dtype=dt)
                                     arr = arr.reshape(entry.tensor_shape)
                                     info_dict[k] = torch.from_numpy(arr)
                                 else:
                                     info_dict[k] = entry.list_data
                     if info_dict:
-                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
+                        setattr(
+                            self.requests[req_id],
+                            "additional_information_cpu",
+                            info_dict,
+                        )
             except Exception:
                 pass
 
@@ -168,33 +177,30 @@ class OmniGPUModelRunner(GPUModelRunner):
                 use_audio_in_video = False
                 for mm_input in self.requests[req_id].mm_inputs:
                     if mm_input.get("image_grid_thw") is not None:
-                        image_grid_thw.extend(
-                            mm_input["image_grid_thw"].tolist())
+                        image_grid_thw.extend(mm_input["image_grid_thw"].tolist())
                     if mm_input.get("video_grid_thw") is not None:
-                        video_grid_thw.extend(
-                            mm_input["video_grid_thw"].tolist())
+                        video_grid_thw.extend(mm_input["video_grid_thw"].tolist())
                     if mm_input.get("second_per_grid_ts") is not None:
-                        second_per_grid_ts.extend(
-                            mm_input["second_per_grid_ts"])
+                        second_per_grid_ts.extend(mm_input["second_per_grid_ts"])
                     if mm_input.get("audio_feature_lengths") is not None:
-                        audio_feature_lengths.extend(
-                            mm_input["audio_feature_lengths"])
+                        audio_feature_lengths.extend(mm_input["audio_feature_lengths"])
                     if mm_input.get("use_audio_in_video") is True:
                         use_audio_in_video = True
 
                 hf_config = self.model_config.hf_config
 
-                self.requests[req_id].mrope_positions, \
-                    self.requests[req_id].mrope_position_delta = \
-                    MRotaryEmbedding.get_input_positions_tensor(
-                        self.requests[req_id].prompt_token_ids,
-                        hf_config=hf_config,
-                        image_grid_thw=image_grid_thw,
-                        video_grid_thw=video_grid_thw,
-                        second_per_grid_ts=second_per_grid_ts,
-                        audio_feature_lengths=audio_feature_lengths,
-                        use_audio_in_video=use_audio_in_video,
-                    )
+                (
+                    self.requests[req_id].mrope_positions,
+                    self.requests[req_id].mrope_position_delta,
+                ) = MRotaryEmbedding.get_input_positions_tensor(
+                    self.requests[req_id].prompt_token_ids,
+                    hf_config=hf_config,
+                    image_grid_thw=image_grid_thw,
+                    video_grid_thw=video_grid_thw,
+                    second_per_grid_ts=second_per_grid_ts,
+                    audio_feature_lengths=audio_feature_lengths,
+                    use_audio_in_video=use_audio_in_video,
+                )
 
             req_ids_to_add.append(req_id)
 
@@ -217,20 +223,19 @@ class OmniGPUModelRunner(GPUModelRunner):
                 new_token_ids = req_data.new_token_ids[i]
                 # Add the sampled token(s) from the previous step (if any).
                 # This doesn't include "unverified" tokens like spec tokens.
-                num_new_tokens = (num_computed_tokens + len(new_token_ids) -
-                                  req_state.num_tokens)
+                num_new_tokens = (
+                    num_computed_tokens + len(new_token_ids) - req_state.num_tokens
+                )
                 if num_new_tokens == 1:
                     # Avoid slicing list in most common case.
                     req_state.output_token_ids.append(new_token_ids[-1])
                 elif num_new_tokens > 0:
-                    req_state.output_token_ids.extend(
-                        new_token_ids[-num_new_tokens:])
+                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
 
             # Update the block IDs.
             if not resumed_from_preemption:
                 # Append the new blocks to the existing block IDs.
-                for block_ids, new_ids in zip(req_state.block_ids,
-                                              new_block_ids):
+                for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
                     block_ids.extend(new_ids)
             else:
                 # The request is resumed from preemption.
@@ -246,8 +251,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 continue
 
             # Update the persistent batch.
-            self.input_batch.num_computed_tokens_cpu[req_index] = (
-                num_computed_tokens)
+            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
             self.input_batch.block_table.append_row(new_block_ids, req_index)
 
             # For the last rank, we don't need to update the token_ids_cpu
@@ -257,21 +261,22 @@ class OmniGPUModelRunner(GPUModelRunner):
                 start_token_index = num_computed_tokens
                 end_token_index = num_computed_tokens + len(new_token_ids)
                 self.input_batch.token_ids_cpu[
-                    req_index,
-                    start_token_index:end_token_index] = new_token_ids
-                self.input_batch.num_tokens_no_spec[
-                    req_index] = end_token_index
+                    req_index, start_token_index:end_token_index
+                ] = new_token_ids
+                self.input_batch.num_tokens_no_spec[req_index] = end_token_index
                 self.input_batch.num_tokens[req_index] = end_token_index
 
             # Add spec_token_ids to token_ids_cpu.
-            spec_token_ids = (
-                scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
+            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
+                req_id, ()
+            )
             if spec_token_ids:
                 num_spec_tokens = len(spec_token_ids)
                 start_index = self.input_batch.num_tokens_no_spec[req_index]
                 end_token_index = start_index + num_spec_tokens
                 self.input_batch.token_ids_cpu[
-                    req_index, start_index:end_token_index] = spec_token_ids
+                    req_index, start_index:end_token_index
+                ] = spec_token_ids
                 # NOTE(woosuk): `num_tokens` here may include spec tokens.
                 self.input_batch.num_tokens[req_index] += num_spec_tokens
 
@@ -287,7 +292,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         self._may_reorder_batch(scheduler_output)
         # Refresh batch metadata with any pending updates.
         self.input_batch.refresh_metadata()
-    
+
     @torch.inference_mode()
     def execute_model(
         self,
@@ -300,29 +305,35 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # Return empty ModelRunnerOutput if there's no work to do.
                 return EMPTY_MODEL_RUNNER_OUTPUT
 
-            return self.kv_connector_no_forward(scheduler_output,
-                                                self.vllm_config)
+            return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
 
         # Prepare the decoder inputs.
-        (attn_metadata, attention_cuda_graphs, logits_indices,
-         spec_decode_metadata, num_scheduled_tokens_np,
-         spec_decode_common_attn_metadata) = (
-             self._prepare_inputs(scheduler_output))
+        (
+            attn_metadata,
+            attention_cuda_graphs,
+            logits_indices,
+            spec_decode_metadata,
+            num_scheduled_tokens_np,
+            spec_decode_common_attn_metadata,
+        ) = self._prepare_inputs(scheduler_output)
 
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if (self.use_cuda_graph
-                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
+        if (
+            self.use_cuda_graph
+            and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]
+        ):
             # Use piecewise CUDA graphs.
             # Add padding to the batch size.
-            num_input_tokens = self.vllm_config.pad_for_cudagraph(
-                num_scheduled_tokens)
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)
         else:
             # Eager mode.
             # Pad tokens to multiple of tensor_parallel_size when
             # enabled collective fusion for SP
             tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.compilation_config.pass_config. \
-                enable_sequence_parallelism and tp_size > 1:
+            if (
+                self.compilation_config.pass_config.enable_sequence_parallelism
+                and tp_size > 1
+            ):
                 num_input_tokens = round_up(num_scheduled_tokens, tp_size)
             else:
                 num_input_tokens = num_scheduled_tokens
@@ -345,13 +356,13 @@ class OmniGPUModelRunner(GPUModelRunner):
             # multimodal and text-only models can consume embeddings.
             inputs_embeds_scheduled = self.model.get_input_embeddings(
                 input_ids=self.input_ids[:num_scheduled_tokens],
-                multimodal_embeddings=(mm_embeds or None)
-                if self.is_multimodal_model else None,
+                multimodal_embeddings=(
+                    (mm_embeds or None) if self.is_multimodal_model else None
+                ),
             )
 
             # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds[:num_scheduled_tokens].copy_(
-                inputs_embeds_scheduled)
+            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
 
             # Overlay per-request custom prompt_embeds for the prompt portion
             for req_index, req_id in enumerate(self.input_batch.req_ids):
@@ -360,26 +371,27 @@ class OmniGPUModelRunner(GPUModelRunner):
                 if pe_cpu is None:
                     continue
                 num_computed_tokens = int(
-                    self.input_batch.num_computed_tokens_cpu[req_index])
+                    self.input_batch.num_computed_tokens_cpu[req_index]
+                )
                 prompt_len = len(req_state.prompt_token_ids)
                 prompt_remaining = max(0, prompt_len - num_computed_tokens)
                 sched_tokens = int(num_scheduled_tokens_np[req_index])
                 overlay_len = min(sched_tokens, prompt_remaining)
                 if overlay_len <= 0:
                     continue
-                src = pe_cpu[num_computed_tokens:
-                             num_computed_tokens + overlay_len].to(
-                                 dtype=self.dtype,
-                                 device=self.device,
-                                 non_blocking=True)
+                src = pe_cpu[
+                    num_computed_tokens : num_computed_tokens + overlay_len
+                ].to(dtype=self.dtype, device=self.device, non_blocking=True)
                 start_offset = int(self.query_start_loc_cpu[req_index])
-                self.inputs_embeds[start_offset:start_offset + overlay_len]\
-                    .copy_(src)
+                self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
 
             input_ids = self.input_ids[:num_input_tokens]  # preserved for APIs
             inputs_embeds = self.inputs_embeds[:num_input_tokens]
-            model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
-                               if self.is_multimodal_model else {})
+            model_mm_kwargs = (
+                self._extract_mm_kwargs(scheduler_output)
+                if self.is_multimodal_model
+                else {}
+            )
         else:
             # For non-first PP ranks, use token ids as usual; embeddings are
             # only consumed on the first rank.
@@ -395,7 +407,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             intermediate_tensors = None
         else:
             intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True)
+                num_input_tokens, intermediate_tensors, True
+            )
 
         # Some attention backends only support CUDA Graphs in pure decode.
         # If attention doesn't support CUDA Graphs for this batch, but we
@@ -405,13 +418,12 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         with set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_input_tokens,
-                num_tokens_across_dp=num_tokens_across_dp,
-                skip_cuda_graphs=skip_cuda_graphs,
-        ), self.maybe_get_kv_connector_output(
-                scheduler_output) as kv_connector_output:
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=num_input_tokens,
+            num_tokens_across_dp=num_tokens_across_dp,
+            skip_cuda_graphs=skip_cuda_graphs,
+        ), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output:
 
             model_output = self.model(
                 input_ids=input_ids,
@@ -433,37 +445,50 @@ class OmniGPUModelRunner(GPUModelRunner):
             hidden_states = model_output
             aux_hidden_states = None
 
-        text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+        text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            hidden_states
+        )
 
         # Broadcast PP output for external_launcher (torchrun)
         # to make sure we are synced across pp ranks
         # TODO: Support overlapping mirco-batches
         # https://github.com/vllm-project/vllm/issues/18019
-        broadcast_pp_output = \
-            self.parallel_config.distributed_executor_backend \
-            == "external_launcher" and len(get_pp_group().ranks) > 0
+        broadcast_pp_output = (
+            self.parallel_config.distributed_executor_backend == "external_launcher"
+            and len(get_pp_group().ranks) > 0
+        )
         if not get_pp_group().is_last_rank:
             # For mid-pipeline stages, return the hidden states.
             assert isinstance(text_hidden_states, IntermediateTensors)
             if not broadcast_pp_output:
                 text_hidden_states.kv_connector_output = kv_connector_output
                 return text_hidden_states
-            get_pp_group().send_tensor_dict(text_hidden_states.tensors,
-                                            all_gather_group=get_tp_group())
+            get_pp_group().send_tensor_dict(
+                text_hidden_states.tensors, all_gather_group=get_tp_group()
+            )
             logits = None
         else:
             if self.input_batch.pooling_params:
-                return self._pool(text_hidden_states, num_scheduled_tokens,
-                                  num_scheduled_tokens_np, kv_connector_output)
+                return self._pool(
+                    text_hidden_states,
+                    num_scheduled_tokens,
+                    num_scheduled_tokens_np,
+                    kv_connector_output,
+                )
 
             sample_hidden_states = text_hidden_states[logits_indices]
             logits = self.model.compute_logits(sample_hidden_states, None)
         if broadcast_pp_output:
-            model_output_broadcast_data = {
-                "logits": logits.contiguous(),
-            } if logits is not None else {}
+            model_output_broadcast_data = (
+                {
+                    "logits": logits.contiguous(),
+                }
+                if logits is not None
+                else {}
+            )
             model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
-                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
+                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
+            )
             assert model_output_broadcast_data is not None
             logits = model_output_broadcast_data["logits"]
 
@@ -513,8 +538,10 @@ class OmniGPUModelRunner(GPUModelRunner):
         discard_sampled_tokens_req_indices = []
         for i, req_id in enumerate(self.input_batch.req_ids):
             req_state = self.requests[req_id]
-            seq_len = (req_state.num_computed_tokens +
-                       scheduler_output.num_scheduled_tokens[req_id])
+            seq_len = (
+                req_state.num_computed_tokens
+                + scheduler_output.num_scheduled_tokens[req_id]
+            )
             if seq_len < req_state.num_tokens:
                 # Ignore the sampled token for partial prefills.
                 # Rewind the generator state as if the token was not sampled.
@@ -529,8 +556,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         # NOTE: GPU -> CPU Sync happens here.
         # Move as many CPU operations as possible before this sync point.
         logprobs_tensors = sampler_output.logprobs_tensors
-        logprobs_lists = logprobs_tensors.tolists() \
-            if logprobs_tensors is not None else None
+        logprobs_lists = (
+            logprobs_tensors.tolists() if logprobs_tensors is not None else None
+        )
 
         # Compute prompt logprobs if needed.
         prompt_logprobs_dict = self._get_prompt_logprobs_dict(
@@ -539,7 +567,8 @@ class OmniGPUModelRunner(GPUModelRunner):
         )
 
         # Get the valid generated tokens.
-        import os
+        # import os  # unused import
+
         sampled_token_ids = sampler_output.sampled_token_ids
         max_gen_len = sampled_token_ids.shape[-1]
         if max_gen_len == 1:
@@ -569,10 +598,10 @@ class OmniGPUModelRunner(GPUModelRunner):
             assert end_idx <= self.max_model_len, (
                 "Sampled token IDs exceed the max model length. "
                 f"Total number of tokens: {end_idx} > max_model_len: "
-                f"{self.max_model_len}")
+                f"{self.max_model_len}"
+            )
 
-            self.input_batch.token_ids_cpu[req_idx,
-                                           start_idx:end_idx] = sampled_ids
+            self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
             self.input_batch.num_tokens_no_spec[req_idx] = end_idx
             self.input_batch.num_tokens[req_idx] = end_idx
             req_id = self.input_batch.req_ids[req_idx]
@@ -609,10 +638,15 @@ class OmniGPUModelRunner(GPUModelRunner):
             num_nans_in_logits=num_nans_in_logits,
             multimodal_outputs=multimodal_outputs,
         )
-    
+
     @torch.inference_mode()
-    def extract_multimodal_outputs(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, dict]:
-        if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
+    def extract_multimodal_outputs(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, dict]:
+        if (
+            hasattr(self.model, "have_multimodal_outputs")
+            and self.model.have_multimodal_outputs
+        ):
             text_hidden_states = hidden_states.text_hidden_states
             multimodal_outputs = hidden_states.multimodal_outputs
 
@@ -620,7 +654,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             text_hidden_states = hidden_states
             multimodal_outputs = {}
         return text_hidden_states, multimodal_outputs
-    
+
     @torch.inference_mode()
     def _dummy_run(
         self,
@@ -645,8 +679,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         num_scheduled_tokens_list[-1] += num_tokens % num_reqs
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
-        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
-                                        dtype=np.int32)
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
 
         attn_metadata: Optional[dict[str, Any]] = None
         if capture_attn_cudagraph:
@@ -655,36 +688,43 @@ class OmniGPUModelRunner(GPUModelRunner):
             # Make sure max_model_len is used at the graph capture time.
             self.seq_lens_np[:num_reqs] = self.max_model_len
             self.seq_lens_np[num_reqs:] = 0
-            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
-                                           non_blocking=True)
+            self.seq_lens[:num_reqs].copy_(
+                self.seq_lens_cpu[:num_reqs], non_blocking=True
+            )
 
             for kv_cache_group_id, kv_cache_group_spec in enumerate(
-                    self.kv_cache_config.kv_cache_groups):
+                self.kv_cache_config.kv_cache_groups
+            ):
                 common_attn_metadata = CommonAttentionMetadata(
-                    query_start_loc=self.query_start_loc[:num_reqs + 1],
-                    query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
-                                                                 1],
+                    query_start_loc=self.query_start_loc[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
                     seq_lens=self.seq_lens[:num_reqs],
                     seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
-                    num_computed_tokens_cpu=self.input_batch.
-                    num_computed_tokens_cpu_tensor[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[  # noqa: E501
+                        :num_reqs
+                    ],
                     num_reqs=num_reqs,
                     num_actual_tokens=num_tokens,
                     max_query_len=num_tokens,
                     block_table_tensor=self.input_batch.block_table[
-                        kv_cache_group_id].get_device_tensor()[:num_reqs],
-                    slot_mapping=self.input_batch.
-                    block_table[kv_cache_group_id].slot_mapping[:num_tokens],
-                    causal=True)
+                        kv_cache_group_id
+                    ].get_device_tensor()[:num_reqs],
+                    slot_mapping=self.input_batch.block_table[
+                        kv_cache_group_id
+                    ].slot_mapping[:num_tokens],
+                    causal=True,
+                )
 
                 for attn_group in self.attn_groups[kv_cache_group_id]:
-                    attn_metadata_i = attn_group.metadata_builder\
-                        .build_for_cudagraph_capture(common_attn_metadata)
+                    attn_metadata_i = (
+                        attn_group.metadata_builder.build_for_cudagraph_capture(
+                            common_attn_metadata
+                        )
+                    )
                     for layer_name in kv_cache_group_spec.layer_names:
                         attn_metadata[layer_name] = attn_metadata_i
 
-        with self.maybe_dummy_run_with_lora(self.lora_config,
-                                            num_scheduled_tokens):
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens):
             if self.is_multimodal_model:
                 input_ids = None
                 inputs_embeds = self.inputs_embeds[:num_tokens]
@@ -707,16 +747,20 @@ class OmniGPUModelRunner(GPUModelRunner):
                         self.model.make_empty_intermediate_tensors(
                             batch_size=self.max_num_tokens,
                             dtype=self.model_config.dtype,
-                            device=self.device))
+                            device=self.device,
+                        )
+                    )
 
                 intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                    num_tokens, None, False)
+                    num_tokens, None, False
+                )
 
             with self.maybe_randomize_inputs(input_ids), set_forward_context(
-                    attn_metadata,
-                    self.vllm_config,
-                    num_tokens=num_tokens,
-                    num_tokens_across_dp=num_tokens_across_dp):
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+            ):
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,
@@ -726,7 +770,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                         model_mm_kwargs,
                         device=self.device,
                     ),
-                    sampler = None
+                    sampler=None,
                 )
 
             if self.use_aux_hidden_state_outputs:
@@ -734,8 +778,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             else:
                 hidden_states = outputs
 
-            text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
-
+            text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+                hidden_states
+            )
 
             if self.speculative_config and self.speculative_config.use_eagle():
                 assert isinstance(self.drafter, EagleProposer)
@@ -752,4 +797,4 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.eplb_step(is_dummy=True, is_profile=is_profile)
 
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        return text_hidden_states, text_hidden_states[logit_indices]
+        return text_hidden_states, text_hidden_states[logit_indices]

vllm_omni/worker/model_runner.py

@@ -1,14 +1,12 @@
-from dataclasses import dataclass
-from typing import Optional, Dict, Any, List
 import itertools
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
 
 import torch
-
-from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder
-
-from vllm.utils import flatten_2d_lists, async_tensor_h2d
 from vllm.lora.layers import LoRAMapping
 from vllm.sequence import MultiModalKwargs
+from vllm.utils import async_tensor_h2d, flatten_2d_lists
+from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder
 
 
 @dataclass(frozen=True)
@@ -34,8 +32,9 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
             if inter_data.inputs_embeds is not None:
                 inputs_embeds_list.append(
                     inter_data.inputs_embeds.to(
-                        dtype=self.runner.model_config.dtype,
-                        device=self.runner.device))
+                        dtype=self.runner.model_config.dtype, device=self.runner.device
+                    )
+                )
             # Support v1 direct-transfer prompt embeds attached on request state
             if inter_data.inputs_embeds is None:
                 try:
@@ -46,7 +45,8 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
                     if pe_cpu is not None:
                         inter_data.inputs_embeds = pe_cpu.to(
                             dtype=self.runner.model_config.dtype,
-                            device=self.runner.device)
+                            device=self.runner.device,
+                        )
                         inputs_embeds_list.append(inter_data.inputs_embeds)
                 except Exception:
                     pass
@@ -55,8 +55,8 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
             inputs_embeds = None
         else:
             inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to(
-                dtype=self.runner.model_config.dtype,
-                device=self.runner.device)
+                dtype=self.runner.model_config.dtype, device=self.runner.device
+            )
             assert len(inputs_embeds) == len(input_tokens)
 
         if not input_tokens and inputs_embeds is None:
@@ -65,20 +65,22 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
             return self.model_input_cls()
 
         mrope_input_positions: Optional[List[List[int]]] = None
-        if any(inter_data.mrope_input_positions is not None
-               for inter_data in self.inter_data_list):
+        if any(
+            inter_data.mrope_input_positions is not None
+            for inter_data in self.inter_data_list
+        ):
             mrope_input_positions = [[] for _ in range(3)]
             for idx in range(3):
                 for inter_data in self.inter_data_list:
                     msections = inter_data.mrope_input_positions
                     if msections is None:
                         for _seq_input_positions in inter_data.input_positions:
-                            mrope_input_positions[idx].extend(
-                                _seq_input_positions)
+                            mrope_input_positions[idx].extend(_seq_input_positions)
                     else:
                         for _seq_mrope_input_positions in msections:
                             mrope_input_positions[idx].extend(
-                                _seq_mrope_input_positions[idx])
+                                _seq_mrope_input_positions[idx]
+                            )
             input_positions = None
         else:
             input_positions = []
@@ -94,23 +96,23 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
             seq_lens.extend(inter_data.seq_lens)
             query_lens.extend(inter_data.query_lens)
             if not inter_data.is_prompt:
-                max_decode_seq_len = max(max_decode_seq_len,
-                                         max(inter_data.seq_lens))
+                max_decode_seq_len = max(max_decode_seq_len, max(inter_data.seq_lens))
                 if self.runner.model_config.is_encoder_decoder:
-                    max_encoder_seq_len = max(max_encoder_seq_len,
-                                              inter_data.encoder_seq_len)
+                    max_encoder_seq_len = max(
+                        max_encoder_seq_len, inter_data.encoder_seq_len
+                    )
 
         # Mapping from request IDs to sequence IDs. Used for Jamba models
         # that manages the cache by itself.
         request_ids_to_seq_ids = {
-            data.request_id: data.seq_ids
-            for data in self.inter_data_list
+            data.request_id: data.seq_ids for data in self.inter_data_list
         }
 
         cuda_graph_pad_size = self._get_cuda_graph_pad_size(
             num_seqs=len(seq_lens),
             max_decode_seq_len=max_decode_seq_len,
-            max_encoder_seq_len=max_encoder_seq_len)
+            max_encoder_seq_len=max_encoder_seq_len,
+        )
 
         batch_size = len(input_tokens)
         if cuda_graph_pad_size != -1:
@@ -123,63 +125,77 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
         if cuda_graph_pad_size:
             input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
         assert self.runner.device is not None
-        input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,
-                                               self.runner.device,
-                                               self.runner.pin_memory)
+        input_tokens_tensor = async_tensor_h2d(
+            input_tokens, torch.long, self.runner.device, self.runner.pin_memory
+        )
 
-        token_types_tensor = async_tensor_h2d(token_types, torch.long,
-                                               self.runner.device,
-                                               self.runner.pin_memory) \
-                                                if token_types else None
+        token_types_tensor = (
+            async_tensor_h2d(
+                token_types, torch.long, self.runner.device, self.runner.pin_memory
+            )
+            if token_types
+            else None
+        )
 
         if mrope_input_positions is not None:
             for idx in range(3):
                 mrope_input_positions[idx].extend(
-                    itertools.repeat(0, cuda_graph_pad_size))
-            input_positions_tensor = async_tensor_h2d(mrope_input_positions,
-                                                      torch.long,
-                                                      self.runner.device,
-                                                      self.runner.pin_memory)
+                    itertools.repeat(0, cuda_graph_pad_size)
+                )
+            input_positions_tensor = async_tensor_h2d(
+                mrope_input_positions,
+                torch.long,
+                self.runner.device,
+                self.runner.pin_memory,
+            )
         else:
             input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))
-            input_positions_tensor = async_tensor_h2d(input_positions,
-                                                      torch.long,
-                                                      self.runner.device,
-                                                      self.runner.pin_memory)
+            input_positions_tensor = async_tensor_h2d(
+                input_positions, torch.long, self.runner.device, self.runner.pin_memory
+            )
         # Sequence and query lengths.
         if cuda_graph_pad_size:
             seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))
 
         # Attention metadata.
         attn_metadata = self.attn_metadata_builder.build(
-            seq_lens, query_lens, cuda_graph_pad_size, batch_size)
+            seq_lens, query_lens, cuda_graph_pad_size, batch_size
+        )
 
         # LoRA data.
         lora_requests = set()
         lora_mapping = None
         if self.enable_lora:
-            lora_requests = set(r for data in self.inter_data_list
-                                for r in data.lora_requests)
-            lora_index_mapping = flatten_2d_lists([
-                flatten_2d_lists(inter_data.lora_index_mapping)
-                for inter_data in self.inter_data_list
-            ])
+            lora_requests = set(
+                r for data in self.inter_data_list for r in data.lora_requests
+            )
+            lora_index_mapping = flatten_2d_lists(
+                [
+                    flatten_2d_lists(inter_data.lora_index_mapping)
+                    for inter_data in self.inter_data_list
+                ]
+            )
             if cuda_graph_pad_size:
-                lora_index_mapping.extend(
-                    itertools.repeat(0, cuda_graph_pad_size))
-            lora_prompt_mapping = flatten_2d_lists([
-                flatten_2d_lists(inter_data.lora_prompt_mapping)
-                for inter_data in self.inter_data_list
-            ])
+                lora_index_mapping.extend(itertools.repeat(0, cuda_graph_pad_size))
+            lora_prompt_mapping = flatten_2d_lists(
+                [
+                    flatten_2d_lists(inter_data.lora_prompt_mapping)
+                    for inter_data in self.inter_data_list
+                ]
+            )
 
             lora_mapping = LoRAMapping(
-                **dict(index_mapping=lora_index_mapping,
-                       prompt_mapping=lora_prompt_mapping,
-                       is_prefill=not self.decode_only))
+                **dict(
+                    index_mapping=lora_index_mapping,
+                    prompt_mapping=lora_prompt_mapping,
+                    is_prefill=not self.decode_only,
+                )
+            )
 
         # Multi-modal data.
         multi_modal_kwargs_list = [
-            data.multi_modal_kwargs for data in self.inter_data_list
+            data.multi_modal_kwargs
+            for data in self.inter_data_list
             if data.multi_modal_kwargs is not None
         ]
         multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
@@ -196,4 +212,5 @@ class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
             lora_requests=lora_requests,
             multi_modal_kwargs=multi_modal_kwargs,
             request_ids_to_seq_ids=request_ids_to_seq_ids,
-            finished_requests_ids=self.finished_requests_ids)
+            finished_requests_ids=self.finished_requests_ids,
+        )