anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge pull request #36 from Gaohan123/remove_unused

Browse Source 
Remove unused files and fix some bugs
This commit is contained in:
Hongsheng Liu authored and GitHub committed 2025-10-25 23:21:52 +08:00
2 parents 6248bbf646 c19e4d26d1
commit 0b56cd3d5b
21 files changed
+556 -1891
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+4
View File
@@ -51,6 +51,10 @@ cd vllm
git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941
VLLM_USE_PRECOMPILED=1 uv pip install --editable .
```
Install additional requirements for vllm-omni
```bash
uv pip install -r requirements.txt
```
## Run examples (Qwen2.5-omni)
docs/architecture/implementation_architecture.md
-353
View File
@@ -1,353 +0,0 @@
# vLLM-omni Implementation Architecture
## 1. Package Structure
```
vllm_omni/
├── __init__.py # Main package exports
├── config/ # Configuration management
│ ├── __init__.py
│ ├── stage_config.py # OmniStageConfig implementation
│ ├── dit_config.py # DiT-specific configurations
│ └── cache_config.py # Caching configurations
├── core/ # Core processing components
│ ├── __init__.py
│ ├── stage_manager.py # Multi-stage orchestration
│ ├── dit_cache_manager.py # DiT caching system
│ └── sched/ # Schedulers
│ ├── __init__.py
│ ├── scheduler.py # Base scheduler interface
├── engine/ # Engine components
│ ├── __init__.py
│ ├── processor.py # Output processing
│ └── output_processor.py # Multimodal output handling
├── executor/ # Executor implementations
│ ├── __init__.py
│ ├── base_executor.py # Base executor interface
│ └── diffusers_executor.py # Diffusers pipeline executor
├── model_executor/ # Model execution components
│ ├── __init__.py
├── worker/ # Worker implementations
│ ├── __init__.py
│ ├── gpu_diffusion_model_runner.py # Loads and runs diffusers pipelines
│ └── gpu_diffusion_worker.py # Worker facade around the diffusion runner
├── entrypoints/ # Entry points and CLI
│ ├── __init__.py
│ ├── omni.py # OmniServeCommand
│ ├── omni_llm.py # OmniLLM and AsyncOmniLLM
│ └── cli/ # CLI integration
│ ├── __init__.py
│ └── main.py # CLI main entry point
├── request.py # Request handling
├── dit_cache_interface.py # DiT cache interface
└── utils/ # Utility functions
├── __init__.py
├── multimodal.py # Multimodal utilities
└── vae.py # VAE utilities for image processing
```
## 2. Core Module Dependencies
### 2.1 vLLM Integration Points
```python
# Key vLLM imports and extensions
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.core.sched.scheduler import SchedulerInterface
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm.v1.outputs import ModelRunnerOutput
```
### 2.2 Internal Dependencies
```python
# Internal module dependencies
vllm_omni.config → vllm_omni.core
vllm_omni.core → vllm_omni.engine, vllm_omni.executor, vllm_omni.worker
vllm_omni.engine → vllm_omni.model_executor
vllm_omni.entrypoints → vllm_omni.core
```
## 3. Configuration System
### 3.1 Stage Configuration
```python
# vllm_omni/config/stage_config.py
@dataclass
class OmniStageConfig:
stage_id: int
engine_type: Literal["AR", "DiT"]
model_path: str
input_modalities: List[str]
output_modalities: List[str]
vllm_config: Optional[VllmConfig] = None
dit_config: Optional[DiTConfig] = None
executor_class: Optional[type[Executor]] = None
stage_output: Optional[Any] = None
```
### 3.2 DiT Configuration
```python
# vllm_omni/config/stage_config.py
@dataclass
class DiTConfig:
model_config: Optional[ModelConfig] = None
scheduler_config: Optional[Any] = None
device_config: Optional[DeviceConfig] = None
load_config: Optional[LoadConfig] = None
compilation_config: Optional[CompilationConfig] = None
dit_cache_config: Optional[DiTCacheConfig] = None
num_inference_steps: int
guidance_scale: float = 7.5
use_diffusers: bool = False
diffusers_pipeline: Optional[str] = None
height: int = 512
width: int = 512
```
### 3.3 Cache Configuration
```python
# vllm_omni/config/cache_config.py
@dataclass
class DiTCacheConfig:
cache_tensors: List[DiTCacheTensor]
max_cache_size: int
cache_strategy: str = "fifo"
enable_optimization: bool = True
```
## 4. Core Implementation Details
### 4.1 OmniLLM Implementation
```python
# vllm_omni/core/omni_llm.py
class OmniLLM(LLM):
def __init__(self, stage_configs: List[OmniStageConfig]):
super().__init__()
self.stage_configs = stage_configs
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:
if stage_config.engine_type == "AR":
engine = self._create_ar_engine(stage_config)
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)
```
### 4.2 AsyncOmniLLM Implementation
```python
# vllm_omni/core/omni_llm.py
class AsyncOmniLLM(AsyncLLM):
def __init__(self, stage_configs: List[OmniStageConfig]):
super().__init__()
self.stage_configs = stage_configs
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
pass
``
### 4.3 Model Runners
```python
# vllm_omni/model_executor/dit_model_runner.py
class OmniDiffusionModelRunner(GPUModelRunner):
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
# DiT model execution logic
return ModelRunnerOutput(
req_ids=[...],
req_id_to_index={...},
sampled_token_ids=[],
spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[tensor, ...], # DiT output tensors
kv_connector_output=None,
num_nans_in_logits=None,
)
# vllm_omni/model_executor/ar_model_runner.py
class OmniARModelRunner(GPUModelRunner):
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
# AR model execution with hidden state output
return ModelRunnerOutput(
req_ids=[...],
req_id_to_index={...},
sampled_token_ids=[],
spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[hidden_states, ...], # Hidden states
kv_connector_output=None,
num_nans_in_logits=None,
)
```
### 4.5 Output Processing
```python
# vllm_omni/engine/output_processor.py
class MultimodalOutputProcessor(OutputProcessor):
def __init__(self):
self.output_handlers: Dict[str, Callable] = {
"image": self._process_image_output,
"text+image": self._process_text_image_output,
"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:
output_type = self._detect_output_type(engine_core_output)
handler = self.output_handlers.get(output_type, self._process_pooling_output)
handler(engine_core_output)
```
## 5. CLI Integration
### 5.1 Entry Point Override
```python
# vllm_omni/entrypoints/cli/main.py
def main():
"""Main CLI entry point that intercepts vLLM commands"""
if "--omni" in sys.argv:
omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"]
omni_serve = OmniServeCommand()
omni_serve.run(omni_args)
else:
from vllm.entrypoints.cli.main import main as vllm_main
vllm_main()
```
### 5.2 Package Configuration
```toml
# pyproject.toml updates
[project.scripts]
vllm = "vllm_omni.entrypoints.cli.main:main"
vllm-omni = "vllm_omni.entrypoints.cli.main:main"
[project.entry-points."vllm.plugins"]
omni = "vllm_omni.plugin:OmniPlugin"
```
## 6. Testing Strategy
### 6.1 Unit Tests
- Individual component testing
- Mock vLLM dependencies
- Configuration validation
### 6.2 Integration Tests
- End-to-end pipeline testing
- vLLM compatibility testing
- Multi-stage processing validation
### 6.3 Performance Tests
- Benchmarking against native vLLM
- Memory usage profiling
- Latency measurements
## 7. Installation and Setup
### 7.1 Package Installation
```bash
pip install vllm>=0.10.2
pip install vllm-omni
```
### 7.2 Development Setup
```bash
git clone https://github.com/hsliuustc0106/vllm-omni
cd vllm-omni
pip install -e ".[dev]"
```
### 7.3 Usage
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
```
## 8. Memory Management
### 8.1 DiT Cache Management
- Tensor caching for intermediate results
- Memory pooling for efficient allocation
- Cache eviction strategies
### 8.2 Multi-Stage Memory
- Inter-stage data passing optimization
- Memory sharing between stages
- Garbage collection optimization
## 9. Error Handling
### 9.1 Stage Failure Handling
- Graceful degradation on stage failures
- Error propagation and reporting
- Recovery mechanisms
### 9.2 vLLM Compatibility
- Version compatibility checks
- API change detection
- Fallback mechanisms
## 10. Monitoring and Logging
### 10.1 Performance Metrics
- Stage execution times
- Memory usage per stage
- Throughput measurements
### 10.2 Debug Information
- Request tracing across stages
- Cache hit/miss ratios
- Error logging and reporting
docs/architecture/vllm_omni_design.md
+509
View File
@@ -0,0 +1,509 @@
# vLLM-omni Software Design Document
## Overview
vLLM-omni is a multi-modality extension for vLLM that supports non-autoregressive structures and non-textual outputs. This document outlines the key software abstractions, APIs, and dependencies for the system, designed to maximize reuse of vLLM's proven architecture.
## Architecture Principles
1. **vLLM V1 Compatibility**: Built on vLLM's Engine V1 architecture with AsyncLLM and EngineCore patterns
2. **Stage-based Processing**: Models are divided into multiple stages, each processed by different Engine Cores
3. **Multiple Engine Core Support**: Each stage can use either AR Engine Core (reusing vLLM) or Diffusion Engine Core (new DiT support), or other new Cores
4. **Worker Process Pattern**: Follows vLLM's multiprocess worker architecture for scalability
5. **Extensibility**: Easy integration of new modalities, model architectures, and output formats
## Key Data Flow
API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
--> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
--> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)
## Core Components (Ordered by Data Flow)
### 1. Installation
```bash
pip install vllm
pip install vllm-omni
```
### 2. Online inference launch (Entry Point)
Keep consistency with the vllm main branch
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
```
#### Design Architecture
```mermaid
graph TD
A[vllm serve model --omni] --> B[vLLM-omni CLI Wrapper]
B --> C{Detect --omni flag}
C -->|Yes| D[Parse OmniConfig]
C -->|No| E[Forward to vLLM CLI]
D --> F[Initialize AsyncOmniLLM]
F --> G[Start omni Server]
G --> H[Multi-stage Processing]
E --> I[Standard vLLM Pipeline]
subgraph "vLLM-omni Components"
F
G
H
end
subgraph "vLLM Components"
I
end
```
#### omni Serve Command Implementation
File: vllm_omni/entrypoints/cli/main.py
```python
import sys
import argparse
from typing import List, Optional
from vllm_omni.entrypoints.omni import OmniServeCommand
def main():
"""Main CLI entry point that intercepts vLLM commands"""
# Check if --omni flag is present
if "--omni" in sys.argv:
# Remove --omni flag and process with vLLM-omni
omni_args = [arg for arg in sys.argv[1:] if arg != "--omni"]
omni_serve = OmniServeCommand()
omni_serve.run(omni_args)
else:
# Forward to original vLLM CLI
from vllm.entrypoints.cli.main import main as vllm_main
vllm_main()
if __name__ == "__main__":
main()
```
#### vLLM Plugin Integration
File: vllm_omni/pyproject.toml
```python
[project.scripts]
# Override vLLM's CLI entry point
vllm = "vllm_omni.cli:main"
# Add entry point for vLLM integration
[project.entry-points."vllm.plugins"]
omni = "vllm_omni.plugin:OmniPlugin"
```
### 3. Offline inference launch (Entry Point)
Design an upper level class to incorporate multi Engines, each engine has a engine call.
```python
from vllm.entrypoints.llm import LLM
from vllm.v1.engine.llm_engine import LLMEngine
class OmniLLM(LLM):
"""Extended LLM supporting multiple engines and stage-based processing"""
def __init__(self, stage_configs: List[StageConfig]):
super().__init__()
self.stage_configs = stage_configs
self.engine_list = [] # List of AsyncLLM instances for each stage
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:
stage_llm = LLMEngine.from_vllm_config(
vllm_config=stage_config.vllm_config,
executor_class=self.executor_class,
log_stats=self.log_stats
)
self.engine_list.append(stage_llm)
def generate(
self,
stage_args_list: List[stage_args],
use_tqdm: Union[bool, Callable[..., tqdm]] = True,
lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
priority: Optional[list[int]] = None,
) -> list[RequestOutput]:
"""Main generation interface - orchestrates multi-stage processing"""
# Process through each stage sequentially
for i, stage_config in enumerate(self.stage_configs):
stage_engine = self.engine_list[i]
# Prepare input for this stage
stage_args = stage_args_list[i]
prompt_str, engine_request, tokenization_kwargs = self._process_stage_inputs(stage_config, **stage_args)
# Add inputs to Engine
stage_engine.add_request(requesy_id, prompt_str, tokenization_kwargs)
# Run Engine
stage_output = stage_engine.step()
# Update stage input and output of each stage for later usage
self._update_stage_io(stage_output, stage_config)
# Process final output
final_output = self.output_processor.process_output(
self.stage_configs[-1].stage_output
)
return final_output
def _process_stage_inputs(stage_config, **stage_args) -> Any:
"""Prepare input for specific stage"""
if stage_config.engine_type == "AR":
return self._process_ar_inputs(**stage_args)
elif stage_config.engine_type == "DiT":
return self._process_dit_inputs(**stage_args)
else
raise NotImplementedError
def _process_dit_inputs(**stage_args)-> Any:
image_latent = self.vae.encode(**stage_args)
return image_latent
```
### 4. Stage config for setting of different model stages
```python
@dataclass
class StageConfig:
"""Configuration for a processing stage"""
stage_id: int
engine_type: str # "AR" or "DiT"
model_path: str
input_modalities: List[str]
output_modalities: List[str]
vllm_config: Optional[VllmConfig] = None # For engine config of corresponding stage
executor_class: type[Executor] # For execute class config of corresponding stage
dit_config: Optional[DiTConfig] = None # For diffusion stages
cache_config: Optional[DiTCacheConfig] = None
```
For AR stage, the setting is:
```python
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
ar_stage_config = StageConfig()
ar_stage_config.vllm_config.scheduler_config.scheduler_cls = Scheduler # original vllm scheduler for AR
ar_stage_config.executor_class = MultiprocExecutor # original vllm executor for AR
```
For DiT stage, the setting is:
```python
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
dit_stage_config = StageConfig()
dit_stage_config.vllm_config.scheduler_config.scheduler_cls = DiffusionScheduler # New scheduler for DiT
# For diffusion models without using diffusers
dit_stage_config.executor_class = MultiprocExecutor
# For diffusion models using diffusers
dit_stage_config.executor_class = DiffusersPipelineExecutor
```
### 5. Online inference main class
Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
```python
from vllm.v1.engine.async_llm improt AsyncLLM
class AsyncOmniLLM(AsyncLLM):
"""Extended AsyncLLM supporting multiple engines and stage-based processing"""
def __init__(self, stage_configs: List[StageConfig]):
super().__init__()
```
### 6. Engine Core
No need to change. The specific executor, scheduler and output type can be transferred into it with new configs.
```python
# class: EngineCore
# EngineCore.step (simplified)
class EngineCore:
def step(self):
scheduler_output = scheduler.schedule()
model_output = executor.execute_model(scheduler_output)
engine_outputs = scheduler.update_from_output(
scheduler_output, model_output
)
return engine_outputs
```
### 7. Scheduler
Create a new Diffusion Scheduler for DiT, which is inherited from original vllm AR scheduler. At first, no complicated strategy.
Just process the request first in first out. Then create a class of DiT Cache Manager for optimization in the future.
```python
from vllm.v1.core.sched.scheduler import Scheduler
class DiffusionScheduler(Scheduler):
def __init__(
self,
vllm_config: VllmConfig,
kv_cache_config: KVCacheConfig,
dit_cache_config: DiTCacheConfig,
structured_output_manager: StructuredOutputManager,
mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
include_finished_set: bool = False,
log_stats: bool = False,
) -> None:
self.dit_cache_manager = DiTCacheManager(dit_cache_config)
```
For DiT Cache Manager (Can refer to xDiT):
```python
class DiTCacheManager:
"""Manages DiT-specific caching"""
def __init__(self, config: DiTCacheConfig):
self.cache_tensors: Dict[str, torch.Tensor] = {}
self.cache_groups: List[DiTCacheTensor] = config.dit_cache_tensors
def allocate_cache(self, request_id: str, size: int) -> torch.Tensor
def get_cache(self, request_id: str) -> Optional[torch.Tensor]
def release_cache(self, request_id: str) -> None
def clear_expired_cache(self) -> None
```
### 8. Executor for DiT without diffusers
No need to change. Just import original AR executor.
### 9. Diffusers Pipeline Executor (no worker)
A single-process executor that directly runs the Diffusers pipeline without spawning workers or using RPC. Interfaces remain identical to `Executor` so the EngineCore loop is unchanged.
#### Function map（Pipeline Executor）
##### 1) Inherited and overridden
```python
from concurrent.futures import Future
from typing import Optional, Union, Callable, Any
import torch
import torch.nn as nn
from vllm.v1.executor.abstract import Executor
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import ModelRunnerOutput
from vllm.tasks import SupportedTask
from diffusers import DiffusionPipeline
class DiffusersPipelineExecutor(Executor):
supports_pp: bool = False # Single-process, no TP/PP/DP
def _init_executor(self) -> None:
# Called by ExecutorBase.__init__
self._failure_callback: Optional[Callable[[], None]] = None
self._device = self._resolve_device()
self._dtype = self._resolve_dtype()
self._pipeline = self._build_pipeline(device=self._device, dtype=self._dtype)
self._profiler = None
self._is_failed = False
self.is_sleeping = False
self.sleeping_tags: set[str] = set()
# major functions to build/run diffusers pipeline
def _build_pipeline(self, device:torch.device, dtype:torch.dtype)->DiffusionPipeline:
model_name = "Qwen/Qwen-Image"
self.pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
self.pipe = pipe.to(device)
def _run_pipeline(self, scheduler_output) -> ModelRunnerOutput: ...
positive_magic = {
"en": ", Ultra HD, 4K, cinematic composition.", # for english prompt
"zh": ", 超清，4K，电影级构图." # for chinese prompt
}
# Generate image
prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
negtive_prompt_embeds = self.pipe.embed_prompt(" ")
width, height = aspect_ratios["16:9"]
image = pipe(
prompt_embeds=prompt_embeds,
negtive_prompt_embeds=negtive_prompt_embeds,
width=width,
height=height,
num_inference_steps=50,
true_cfg_scale=4.0,
generator=torch.Generator(device="cuda").manual_seed(42)
).images[0]
output = self.wrap_image_as_ModelRunnerOutput(image)
return output
# ---- Internal helpers (implementation-specific, not public API) ----
def _resolve_device(self): ...
def _resolve_dtype(self): ...
def _get_model(self) -> nn.Module: ...
def _get_and_process_prompt_embeds(self, scheduler_output, positive_magic):
...
#append the positive_magic to prompt and embed them to prompt embed tensors
# Functions related to workers should either raise NotImplementedError
# or return empty defaults in this no-worker executor.
def collective_rpc(self, method, timeout=None, args=(), kwargs=None) -> list[Any]:
# No workers in pipeline executor
raise NotImplementedError("No workers in DiffusersPipelineExecutor")
def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
return # no-op (pipeline already built in _init_executor)
def register_failure_callback(self, callback):
self._failure_callback = callback
def determine_available_memory(self) -> list[int]: # bytes
# Single device; return [available_bytes]. If CPU-only, return [0].
return [self._determine_available_bytes(self._device)]
```
### 10. Worker
#### Inheritance strategy
Prefer reusing the mature GPU Worker end-to-end. Worker class is selected by configuration (vllm_config). Do not add a new executor-specific worker binding. If customization is needed, override only `init_device` to construct the `DiffusionModelRunner`; all other behaviors (device init details, profiling, sleep/wake, PP/TP comms, execute path) remain from `vllm/v1/worker/gpu_worker.py::Worker`.
Inherited and overridden
```python
# Optional: only if you need to plug a custom DiffusionModelRunner.
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
from vllm.v1.worker.diffusion_model_runner import DiffusionModelRunner
class Worker(GPUWorker):
def init_device(self) -> None:
#those related to device check and init
...
self.model_runner = DiffusionModelRunner(self.vllm_config, ...)
```
### 11. Model Runner
### Function map (Model Runner)
#### 1) Inherited and overridden
Those parts relied to the KV Cache will be omitted if we do not register the model to the vllm config. The engine core will view it as do not require KV Cache, and handle it properly
Reuse `vllm/v1/outputs.py::ModelRunnerOutput`：
- DiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors
- ARModelRunner: Use the `pooler_output=[Tensor,...]` to return hidden states.
```python
from typing import Optional, Union
import torch
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm.v1.outputs import ModelRunnerOutput
class DiffusionModelRunner(GPUModelRunner):
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
...
return ModelRunnerOutput(
req_ids=[...],
req_id_to_index={...},
sampled_token_ids=[],
spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[Tensor,...], # Return Hidden states
kv_connector_output=None,
num_nans_in_logits=None,
)# return multi modal tensors via pooler_output=[Tensor,...]
class ARModelRunner(GPUModelRunner):
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
...
return ModelRunnerOutput(
req_ids=[...],
req_id_to_index={...},
sampled_token_ids=[],
spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[Tensor,...], # Return Hidden states
kv_connector_output=None,
num_nans_in_logits=None,
)# return hidden states via pooler_output=[Tensor,...]
```
### 12. RequestState Role and Usage
```bash
# RequestState serves as the per-request state tracker in OutputProcessor:
# - Maintains request-specific state (tokens, logprobs, detokenizer, etc.)
# - Converts EngineCoreOutput → RequestOutput/PoolingRequestOutput
# - Manages request lifecycle from registration to completion
AsyncLLM.add_request()
↓
OutputProcessor.add_request()
↓
RequestState.from_new_request() → 创建请求状态
AsyncLLM.__init__() / AsyncLLM.generate() / AsyncLLM.encode() → 创建一个Background loop 持续从EngineCore获取输出
↓
OutputProcessor.process_outputs() → 更新状态并处理输出
↓
RequestState.make_request_output() → 转换为最终输出,格式为RequestOutput或者PoolingRequestOutput
↓
RequestOutputCollector.put() → 推送到队列（AsyncLLM）
```
Need to add implementation for an existing method
```python
class RequestState:
def __init__(self, request_id: str, parent_req: Optional[ParentRequest],
request_index: int, lora_name: Optional[str],
output_kind: RequestOutputKind, prompt: Optional[str],
prompt_token_ids: list[int], logprobs_processor: Optional[LogprobsProcessor],
detokenizer: Optional[IncrementalDetokenizer],
max_tokens_param: Optional[int], arrival_time: float,
queue: Optional[RequestOutputCollector], log_stats: bool):
def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput:
"""Create PoolingOutput for multimodal/pooling requests"""
return PoolingOutput(data=pooling_output)
```
### 13. OutputProcessor
For hidden state output, original OutputProcessor already support.
We only need to add new one for final multimodal output.
```python
class MultimodalOutputProcessor(OutputProcessor):
"""Handles multimodal output processing"""
def __init__(self):
self.output_handlers: Dict[str, OutputProcessor] = {}
def process_outputs(self, engine_core_outputs: list[EngineCoreOutput], ...):
for engine_core_output in engine_core_outputs:
# Option 1: Use output_type field (if EngineCoreOutput is extended)
if engine_core_output.output_type == "image":
self._process_image_output(engine_core_output)
elif engine_core_output.output_type == "text+image":
self._process_text_image_output(engine_core_output)
elif engine_core_output.output_type == "latents":
self._process_latents_output(engine_core_output)
elif engine_core_output.output_type == "text":
self._process_text_output(engine_core_output)
else:
# Fallback: use existing pooling_output logic
if engine_core_output.pooling_output is not None:
self._process_pooling_output(engine_core_output)
else:
self._process_text_output(engine_core_output)
```
examples/basic/api_client.py
-121
View File
@@ -1,121 +0,0 @@
#!/usr/bin/env python3
"""
Simple API client example for vLLM-omni server.
This example shows how to interact with the vLLM-omni API server
using HTTP requests.
"""
import json
import requests
def test_health(host="localhost", port=8000):
"""Test if the server is healthy."""
url = f"http://{host}:{port}/health"
try:
response = requests.get(url, timeout=5)
if response.status_code == 200:
print("✅ Server is healthy")
return True
else:
print(f"❌ Server health check failed: {response.status_code}")
return False
except requests.exceptions.RequestException as e:
print(f"❌ Cannot connect to server: {e}")
return False
def generate_text(prompt, host="localhost", port=8000, **kwargs):
"""Generate text using the API server."""
url = f"http://{host}:{port}/generate"
# Default parameters
params = {
"prompts": [prompt],
"max_tokens": kwargs.get("max_tokens", 100),
"temperature": kwargs.get("temperature", 0.7),
"top_p": kwargs.get("top_p", 1.0),
"frequency_penalty": kwargs.get("frequency_penalty", 0.0),
"presence_penalty": kwargs.get("presence_penalty", 0.0),
}
# Add optional parameters
if "stop" in kwargs:
params["stop"] = kwargs["stop"]
try:
response = requests.post(url, json=params, timeout=30)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
print(f"❌ API request failed: {e}")
return None
def main():
"""Main example function."""
print("vLLM-omni API Client Example")
print("=" * 40)
# Test server health
if not test_health():
print("Please start the server first:")
print(" vllm serve Qwen/Qwen3-0.6B --omni --port 8000")
return
# Example 1: Simple generation
print("\n1. Simple text generation:")
result = generate_text("Hello, how are you?", max_tokens=50)
if result:
text = result["outputs"][0]["text"]
print(" Input: Hello, how are you?")
print(f" Output: {text}")
# Example 2: Creative writing
print("\n2. Creative writing:")
result = generate_text(
"Write a short story about a robot learning to paint",
max_tokens=150,
temperature=0.9,
)
if result:
text = result["outputs"][0]["text"]
print(" Input: Write a short story about a robot learning to paint")
print(f" Output: {text}")
# Example 3: Question answering
print("\n3. Question answering:")
result = generate_text(
"What is the capital of France?", max_tokens=30, temperature=0.3
)
if result:
text = result["outputs"][0]["text"]
print(" Input: What is the capital of France?")
print(f" Output: {text}")
# Example 4: Multiple prompts
print("\n4. Multiple prompts:")
prompts = ["Tell me a joke", "What's 2+2?", "Describe the color blue"]
for prompt in prompts:
result = generate_text(prompt, max_tokens=30)
if result:
text = result["outputs"][0]["text"]
print(f" Q: {prompt}")
print(f" A: {text}")
# Example 5: Show full response structure
print("\n5. Full response structure:")
result = generate_text("Hello world", max_tokens=20)
if result:
print(" Full API response:")
print(json.dumps(result, indent=2))
print("\n" + "=" * 40)
print("API client examples completed!")
if __name__ == "__main__":
main()
examples/basic/text_generation.py
-138
View File
@@ -1,138 +0,0 @@
#!/usr/bin/env python3
"""
Simple vLLM-omni usage example.
This example demonstrates how to use vLLM-omni for basic text generation
with a single AR (Autoregressive) stage.
"""
import asyncio
from vllm_omni.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"],
)
# Initialize OmniLLM
omni_llm = OmniLLM([stage_config])
# Prepare stage arguments
stage_args = [
{
"prompt": "Hello, how are you today?",
"max_tokens": 50,
"temperature": 0.7,
}
]
# Generate text
print(f"Input: {stage_args[0]['prompt']}")
outputs = omni_llm.generate(stage_args)
# Display results
for i, output in enumerate(outputs):
print(f"Output {i}:")
if hasattr(output, "outputs") and output.outputs:
for completion in output.outputs:
print(f" Text: {completion.text}")
print(f" Finished: {completion.finish_reason != 'length'}")
print(f" Tokens: {len(completion.token_ids)} tokens")
async def async_example():
"""Asynchronous usage example."""
print("\n=== Asynchronous vLLM-omni Example ===")
# Create a simple AR stage configuration
stage_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"],
)
# Initialize AsyncOmniLLM
omni_llm = AsyncOmniLLM([stage_config])
# Prepare stage arguments
stage_args = [
{
"prompt": "What is artificial intelligence?",
"max_tokens": 100,
"temperature": 0.8,
}
]
# Generate text asynchronously
print(f"Input: {stage_args[0]['prompt']}")
outputs = await omni_llm.generate_async(stage_args)
# Display results
for i, output in enumerate(outputs):
print(f"Output {i}:")
if hasattr(output, "outputs") and output.outputs:
for completion in output.outputs:
print(f" Text: {completion.text}")
print(f" Finished: {completion.finish_reason != 'length'}")
print(f" Tokens: {len(completion.token_ids)} tokens")
def multi_prompt_example():
"""Example with multiple prompts."""
print("\n=== Multi-Prompt Example ===")
# Create stage configuration
stage_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"],
)
# Initialize OmniLLM
omni_llm = OmniLLM([stage_config])
# Multiple prompts
prompts = [
"Tell me a joke",
"What's the weather like?",
"Explain quantum computing",
]
for prompt in prompts:
print(f"\nInput: {prompt}")
stage_args = [{"prompt": prompt, "max_tokens": 30, "temperature": 0.9}]
outputs = omni_llm.generate(stage_args)
if outputs and hasattr(outputs[0], "outputs") and outputs[0].outputs:
completion = outputs[0].outputs[0]
print(f"Output: {completion.text}")
if __name__ == "__main__":
print("vLLM-omni Simple Usage Examples")
print("=" * 40)
# Run synchronous example
sync_example()
# Run asynchronous example
asyncio.run(async_example())
# Run multi-prompt example
multi_prompt_example()
print("\n" + "=" * 40)
print("Examples completed!")
examples/omni/README.md
-31
View File
@@ -1,31 +0,0 @@
# Omni Examples
Examples showcasing multi-stage AR + DiT pipelines powered by vLLM-omni.
## AR ➜ DiT with Diffusers Backend
```bash
# Optional: customise models
export AR_MODEL=Qwen/Qwen3-0.6B
export DIT_MODEL=./models/stable-diffusion-2-1/
# Run with defaults (20 steps @ 512x512)
python examples/omni/ar_dit_diffusers.py
# Faster run
python examples/omni/ar_dit_diffusers.py --steps 14 --height 384 --width 384 --guidance 4.5
```
Arguments:
- `--ar-model`: Override AR stage model (default `Qwen/Qwen3-0.6B`).
- `--dit-model`: Override DiT stage model (default env `DIT_MODEL`, else diffusers repo).
- `--steps`: Number of diffusion steps (default 20).
- `--guidance`: Guidance scale (default 5.0).
- `--height` / `--width`: Output resolution (default 512).
- `--seed`: Optional RNG seed.
- `--prompt`, `--temperature`, `--max-tokens`: Control AR stage generation.
- `--output`: Destination path for the generated image file.
Generated images are saved next to the script as `omni_dit_output.png` by default.
See the script for additional tweaks (scheduler, prompt chaining, etc.).
examples/omni/ar_dit_diffusers.py
-175
View File
@@ -1,175 +0,0 @@
"""Run the AR → DiT (diffusers) pipeline using YAML configuration defaults."""
from __future__ import annotations
import argparse
import os
from pathlib import Path
from typing import Dict, List
import yaml
from vllm_omni import (
DiTConfig,
OmniLLM,
create_ar_stage_config,
create_dit_stage_config,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Run a two-stage AR → DiT pipeline with config-derived defaults."
)
parser.add_argument(
"--config",
type=Path,
default=Path(__file__).resolve().parent / "configs" / "ar_dit_local.yaml",
help="Path to a YAML configuration describing the pipeline stages.",
)
parser.add_argument(
"--prompt",
default="A scenic watercolor painting of a lighthouse at sunset",
help="Prompt passed to the AR stage.",
)
parser.add_argument(
"--negative-prompt",
default=None,
help="Optional negative prompt forwarded to the diffusion stage.",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Optional seed for deterministic diffusion sampling.",
)
parser.add_argument(
"--output",
type=Path,
default=Path("./omni_dit_output.png"),
help="Destination path for the generated image.",
)
return parser.parse_args()
def _normalize_modalities(values: List[str] | None, default: List[str]) -> List[str]:
if not values:
return default
return list(values)
def _load_stage_configs_from_yaml(config_path: Path):
with config_path.open("r", encoding="utf-8") as f:
config = yaml.safe_load(f) or {}
stages = config.get("stages", [])
if not stages:
raise ValueError(f"No stages defined in config: {config_path}")
stage_configs = []
for stage in stages:
stage_id = stage["stage_id"]
engine_type = stage["engine_type"]
default_stage_args = stage.get("default_stage_args")
if engine_type == "AR":
stage_configs.append(
create_ar_stage_config(
stage_id=stage_id,
model_path=stage["model_path"],
input_modalities=_normalize_modalities(
stage.get("input_modalities"), ["text"]
),
output_modalities=_normalize_modalities(
stage.get("output_modalities"), ["text"]
),
default_stage_args=default_stage_args,
)
)
elif engine_type == "DiT":
dit_cfg_dict: Dict = dict(stage.get("dit_config", {}))
if "num_inference_steps" not in dit_cfg_dict:
raise ValueError(
"DiT stage requires 'num_inference_steps' in dit_config"
)
dit_cfg = DiTConfig(**dit_cfg_dict)
stage_configs.append(
create_dit_stage_config(
stage_id=stage_id,
model_path=stage["model_path"],
input_modalities=_normalize_modalities(
stage.get("input_modalities"), ["text"]
),
output_modalities=_normalize_modalities(
stage.get("output_modalities"), ["image"]
),
dit_config=dit_cfg,
default_stage_args=default_stage_args,
)
)
else:
raise ValueError(f"Unsupported engine_type '{engine_type}' in config")
return stage_configs
def _apply_env_overrides(stage_configs):
ar_override = os.environ.get("AR_MODEL")
dit_override = os.environ.get("DIT_MODEL")
for stage_config in stage_configs:
if ar_override and stage_config.engine_type == "AR":
stage_config.model_path = ar_override
if dit_override and stage_config.engine_type == "DiT":
stage_config.model_path = dit_override
def main():
args = parse_args()
stage_configs = _load_stage_configs_from_yaml(args.config)
_apply_env_overrides(stage_configs)
omni = OmniLLM(stage_configs)
stage_overrides: Dict[int, Dict[str, object]] = {}
if args.seed is not None or args.negative_prompt:
for stage_config in stage_configs:
if stage_config.engine_type != "DiT":
continue
override = stage_overrides.setdefault(stage_config.stage_id, {})
if args.seed is not None:
override["seed"] = args.seed
if args.negative_prompt:
override["negative_prompt"] = args.negative_prompt
outputs = omni.generate(
prompt=args.prompt,
stage_overrides=stage_overrides if stage_overrides else None,
)
image = None
for request_output in outputs:
for completion in getattr(request_output, "outputs", []) or []:
candidate = getattr(completion, "image", None)
if candidate is not None:
image = candidate
break
if image is not None:
break
if image is None:
print("No image found in outputs.")
return
out_path = args.output.resolve()
try:
image.save(out_path)
print(f"Saved image to {out_path}")
except Exception:
print("Generated output is not a PIL.Image; skipping save.")
if __name__ == "__main__":
main()
examples/omni/configs/ar_dit_local.yaml
-28
View File
@@ -1,28 +0,0 @@
# Example config for running OmniLLM with local AR + DiT models.
stages:
- stage_id: 0
engine_type: AR
model_path: Qwen/Qwen3-0.6B
input_modalities: [text]
output_modalities: [text]
default_stage_args:
max_tokens: 16
temperature: 0.7
- stage_id: 1
engine_type: DiT
model_path: stabilityai/stable-diffusion-2-1
input_modalities: [text]
output_modalities: [image]
default_stage_args:
height: 512
width: 512
dit_config:
num_inference_steps: 20
guidance_scale: 5.0
use_diffusers: true
diffusers_pipeline: auto
height: 512
width: 512
device_config:
device: mps
dtype: float16
requirements.txt
+17 -27
View File
@@ -1,27 +1,17 @@
# Core dependencies
vllm>=0.10.2
torch>=2.7
transformers>=4.30.0
numpy>=1.21.0
pillow>=8.0.0
opencv-python>=4.5.0
PyYAML>=6.0
# API and serving
fastapi>=0.100.0
uvicorn>=0.20.0
pydantic>=2.0.0
gradio>=3.0.0
# Async and distributed computing
asyncio-mqtt>=0.11.0
ray>=2.0.0
# Development dependencies
pytest>=7.0.0
pytest-asyncio>=0.21.0
pytest-cov>=4.0.0
black>=23.0.0
isort>=5.12.0
flake8>=6.0.0
mypy>=1.0.0
cloudpickle==3.1.1
flash_attn==2.8.3
msgspec==0.19.0
numpy==2.3.4
omegaconf==2.3.0
pydantic==2.12.3
torch==2.8.0
transformers==4.57.0
typing_extensions==4.15.0
decord==0.6.0
librosa==0.11.0
packaging==25.0
Pillow==12.0.0
Requests==2.32.5
resampy==0.4.3
soundfile==0.13.1
torchvision==0.23.0
tests/conftest.py
-86
View File
@@ -1,86 +0,0 @@
"""
Shared fixtures and configuration for vLLM-omni tests.
"""
from unittest.mock import Mock
import pytest
import torch
from vllm_omni.config import (
DiTCacheConfig,
DiTConfig,
create_ar_stage_config,
create_dit_stage_config,
)
@pytest.fixture(scope="session")
def device():
"""Get available device for testing."""
return "cuda" if torch.cuda.is_available() else "cpu"
@pytest.fixture
def sample_ar_stage_config():
"""Sample AR stage configuration for testing."""
return create_ar_stage_config(
stage_id=0,
model_path="test-ar-model",
input_modalities=["text"],
output_modalities=["text"],
)
@pytest.fixture
def sample_dit_stage_config():
"""Sample DiT stage configuration for testing."""
dit_config = DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=10,
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,
)
@pytest.fixture
def sample_stage_configs(sample_ar_stage_config, sample_dit_stage_config):
"""Sample stage configurations for testing."""
return [sample_ar_stage_config, sample_dit_stage_config]
@pytest.fixture
def mock_vllm_config():
"""Mock vLLM configuration."""
config = Mock()
config.model = "test-model"
config.tensor_parallel_size = 1
config.pipeline_parallel_size = 1
return config
@pytest.fixture
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
)
]
return DiTCacheConfig(
cache_tensors=cache_tensors,
max_cache_size=1024 * 1024 * 1024, # 1GB
cache_strategy="fifo",
enable_optimization=True,
)
tests/unit/test_config.py
-160
View File
@@ -1,160 +0,0 @@
"""
Unit tests for configuration modules.
"""
import pytest
from vllm_omni.config import (
DiTCacheConfig,
DiTCacheTensor,
DiTConfig,
OmniStageConfig,
create_ar_stage_config,
create_dit_stage_config,
)
class TestOmniStageConfig:
def test_ar_stage_config_creation(self):
"""Test AR stage configuration creation."""
config = OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-model",
input_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(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,
)
assert config.engine_type == "DiT"
assert config.dit_config is not None
assert config.dit_config.num_inference_steps == 50
def test_invalid_engine_type(self):
"""Test validation of engine type."""
with pytest.raises(ValueError):
OmniStageConfig(
stage_id=0,
engine_type="INVALID",
model_path="test-model",
input_modalities=["text"],
output_modalities=["text"],
)
def test_empty_modalities(self):
"""Test validation of empty modalities."""
with pytest.raises(ValueError):
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-model",
input_modalities=[],
output_modalities=["text"],
)
with pytest.raises(ValueError):
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-model",
input_modalities=["text"],
output_modalities=[],
)
def test_dit_requires_config(self):
"""Test that DiT engine requires dit_config."""
with pytest.raises(ValueError):
OmniStageConfig(
stage_id=1,
engine_type="DiT",
model_path="test-dit-model",
input_modalities=["text"],
output_modalities=["image"],
)
class TestDiTConfig:
def test_dit_config_defaults(self):
"""Test DiT configuration with defaults."""
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(
num_inference_steps=50,
use_diffusers=True,
diffusers_pipeline="stable-diffusion",
)
assert config.use_diffusers is True
assert config.diffusers_pipeline == "stable-diffusion"
class TestDiTCacheConfig:
def test_cache_config_creation(self):
"""Test cache configuration creation."""
cache_tensors = [
DiTCacheTensor(
name="test_tensor",
shape=[1, 512, 512],
dtype="float32",
persistent=True,
)
]
config = DiTCacheConfig(
cache_tensors=cache_tensors,
max_cache_size=1024 * 1024 * 1024,
cache_strategy="fifo",
)
assert len(config.cache_tensors) == 1
assert config.max_cache_size == 1024 * 1024 * 1024
assert config.cache_strategy == "fifo"
assert config.enable_optimization is True
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")
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(num_inference_steps=50)
config = create_dit_stage_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
+5 -21
View File
@@ -12,35 +12,19 @@ Architecture:
processing
"""
__version__ = "0.1.0"
__author__ = "vLLM-omni Team"
__email__ = "hsliuustc@gmail.com"
from . import patch
from .config import (
DiTCacheConfig,
DiTConfig,
OmniStageConfig,
create_ar_stage_config,
create_dit_stage_config,
)
from . import patch # noqa: F401
from .config import OmniModelConfig
# Main entry points
from .entrypoints.omni_llm import OmniLLM
from .entrypoints.omni_llm import OmniLLM
__all__ = [
# Version info
"__version__",
"__author__",
"__email__",
# Main components
"OmniLLM",
# AsyncOmniLLM is not supported yet
# "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented
# Configuration
"OmniStageConfig",
"DiTConfig",
"DiTCacheConfig",
"create_ar_stage_config",
"create_dit_stage_config",
"OmniModelConfig",
# All other components are available through their respective modules
# processors.*, schedulers.*, executors.*, etc.
]
vllm_omni/config/__init__.py
-15
View File
@@ -10,15 +10,6 @@ from vllm.config import ModelConfig, config
import vllm_omni.model_executor.models as me_models
from .stage_config import (
DiTCacheConfig,
DiTCacheTensor,
DiTConfig,
OmniStageConfig,
create_ar_stage_config,
create_dit_stage_config,
)
@config
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
@@ -41,10 +32,4 @@ class OmniModelConfig(ModelConfig):
__all__ = [
"OmniModelConfig",
"OmniStageConfig",
"DiTConfig",
"DiTCacheConfig",
"DiTCacheTensor",
"create_ar_stage_config",
"create_dit_stage_config",
]
vllm_omni/config/stage_config.py
-179
View File
@@ -1,179 +0,0 @@
"""
Stage configuration for vLLM-omni multi-stage processing.
"""
from dataclasses import dataclass
from typing import Any, Dict, List, Literal, Optional, Type
from vllm.config import (
CompilationConfig,
DeviceConfig,
LoadConfig,
ModelConfig,
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"
persistent: bool = True
@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
enable_optimization: bool = True
cache_compression: bool = False
@dataclass
class DiTConfig:
"""Configuration for DiT (Diffusion Transformer) stages."""
num_inference_steps: int
"""Number of diffusion inference steps."""
model_config: Optional[ModelConfig] = None
"""Model configuration."""
scheduler_config: Optional[Any] = None
"""Scheduler configuration."""
device_config: Optional[DeviceConfig] = None
"""Device configuration."""
load_config: Optional[LoadConfig] = None
"""Model loading configuration."""
compilation_config: Optional[CompilationConfig] = None
"""Compilation configuration."""
dit_cache_config: Optional[DiTCacheConfig] = None
"""DiT cache configuration."""
guidance_scale: float = 7.5
use_diffusers: bool = False
diffusers_pipeline: Optional[str] = None
height: int = 512
width: int = 512
batch_size: int = 1
@dataclass
class OmniStageConfig:
"""Configuration for a processing stage in vLLM-omni."""
stage_id: int
engine_type: Literal["AR", "DiT"]
model_path: str
input_modalities: List[str]
output_modalities: List[str]
vllm_config: Optional[VllmConfig] = None
dit_config: Optional[DiTConfig] = None
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'"
)
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,
)
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}"
)
def create_ar_stage_config(
stage_id: int,
model_path: str,
input_modalities: List[str] = None,
output_modalities: List[str] = None,
vllm_config: Optional[VllmConfig] = None,
executor_class: Optional[Type[Executor]] = None,
default_stage_args: Optional[Dict[str, Any]] = None,
) -> OmniStageConfig:
"""Create a configuration for an AR (Autoregressive) stage."""
if input_modalities is None:
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(
stage_id=stage_id,
engine_type="AR",
model_path=model_path,
input_modalities=input_modalities,
output_modalities=output_modalities,
vllm_config=vllm_config,
executor_class=executor_class,
default_stage_args=default_stage_args,
)
def create_dit_stage_config(
stage_id: int,
model_path: str,
input_modalities: Optional[List[str]] = None,
output_modalities: Optional[List[str]] = None,
dit_config: Optional[DiTConfig] = None,
executor_class: Optional[Type[Executor]] = None,
default_stage_args: Optional[Dict[str, Any]] = None,
) -> OmniStageConfig:
"""Create a configuration for a DiT (Diffusion Transformer) stage."""
if input_modalities is None:
input_modalities = ["text"]
if output_modalities is None:
output_modalities = ["image"]
if dit_config is None:
dit_config = DiTConfig(num_inference_steps=50)
return OmniStageConfig(
stage_id=stage_id,
engine_type="DiT",
model_path=model_path,
input_modalities=input_modalities,
output_modalities=output_modalities,
dit_config=dit_config,
executor_class=executor_class,
default_stage_args=default_stage_args,
)
vllm_omni/core/dit_cache_manager.py
+4 -176
View File
@@ -1,185 +1,13 @@
"""
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 ..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] = {}
self.cache_groups: List[DiTCacheTensor] = config.cache_tensors
self.max_cache_size = config.max_cache_size
self.current_cache_size = 0
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
)
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:
cache_tensor = self.cache_tensors[request_id]
cache_tensor.access_count += 1
self.cache_hits += 1
return cache_tensor.tensor
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,
)
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,
)
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":
self._evict_fifo()
elif self.config.cache_strategy == "lru":
self._evict_lru()
elif self.config.cache_strategy == "lfu":
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,
)
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,
)
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,
)
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,
"cache_utilization": self.current_cache_size / self.max_cache_size,
"cache_hits": self.cache_hits,
"cache_misses": self.cache_misses,
"hit_rate": hit_rate,
"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
def __init__(self, config):
pass
vllm_omni/engine/diffusion_engine.py
-126
View File
@@ -1,126 +0,0 @@
"""Diffusion engine adapters for diffusers-backed DiT stages."""
from __future__ import annotations
from typing import Any, Optional
from ..config import DiTConfig
from ..worker.gpu_diffusion_worker import DiffusionGPUWorker, DiffusionRunnerOutput
class DiffusionEngine:
"""Base diffusion engine storing shared DiT configuration."""
def __init__(
self,
dit_config: DiTConfig,
model_path: Optional[str] = None,
log_stats: bool = False,
multiprocess_mode: bool = False,
) -> None:
self.dit_config = dit_config
model_cfg = getattr(dit_config, "model_config", None)
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"
)
else:
config_model_path = getattr(model_cfg, "model", None) or getattr(
model_cfg, "model_path", None
)
self.model_path = model_path or config_model_path
self.model_config = dit_config.model_config
self.cache_config = dit_config.dit_cache_config
self.height = dit_config.height
self.width = dit_config.width
self.num_inference_steps = dit_config.num_inference_steps
self.guidance_scale = dit_config.guidance_scale
self.log_stats = log_stats
self.multiprocess_mode = multiprocess_mode
def generate(self, *args, **kwargs): # pragma: no cover - placeholder
raise NotImplementedError("This method should be implemented in subclasses.")
class DiffusersPipelineEngine(DiffusionEngine):
"""Adapter that invokes a diffusers pipeline via DiffusionGPUWorker."""
def __init__(
self,
dit_config: DiTConfig,
model_path: Optional[str] = None,
log_stats: bool = False,
multiprocess_mode: bool = False,
) -> None:
super().__init__(dit_config, model_path, log_stats, multiprocess_mode)
resolved_model_path = self.model_path
if resolved_model_path is None:
raise ValueError("Model path must be provided for DiffusersPipelineEngine.")
device_cfg = getattr(dit_config, "device_config", None)
model_cfg = getattr(dit_config, "model_config", None)
device = None
dtype = None
if device_cfg:
if isinstance(device_cfg, dict):
device = device_cfg.get("device")
dtype = device_cfg.get("dtype")
else:
device = getattr(device_cfg, "device", None)
dtype = getattr(device_cfg, "dtype", None)
if dtype is None and model_cfg:
if isinstance(model_cfg, dict):
dtype = model_cfg.get("dtype")
else:
dtype = getattr(model_cfg, "dtype", None)
self.worker = DiffusionGPUWorker(
resolved_model_path,
pipeline_name=dit_config.diffusers_pipeline,
device=device,
dtype=dtype,
)
def generate(
self,
prompt: str,
*,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: Optional[int] = None,
guidance_scale: Optional[float] = None,
negative_prompt: Optional[str] = None,
seed: Optional[int] = None,
image: Optional[Any] = None,
) -> DiffusionRunnerOutput:
height = int(height) if height is not None else self.height
width = int(width) if width is not None else self.width
num_steps = (
int(num_inference_steps)
if num_inference_steps is not None
else self.num_inference_steps
)
guidance = (
float(guidance_scale) if guidance_scale is not None else self.guidance_scale
)
return self.worker.generate(
prompt=prompt,
height=height,
width=width,
num_inference_steps=num_steps,
guidance_scale=guidance,
negative_prompt=negative_prompt,
seed=seed,
image=image,
)
vllm_omni/entrypoints/api_server.py
-150
View File
@@ -1,150 +0,0 @@
"""
API server for vLLM-omni.
"""
import asyncio
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
stage_args: List[Dict[str, Any]] = None
class GenerateResponse(BaseModel):
"""Response model for generation."""
outputs: List[Dict[str, Any]]
stage_outputs: List[Dict[str, Any]] = None
app = FastAPI(
title="vLLM-omni API",
description="Multi-modality models inference and serving",
version="0.1.0",
)
# Global omni_llm instance
omni_llm: AsyncOmniLLM = None
@app.on_event("startup")
async def startup_event():
"""Initialize the omni_llm instance on startup."""
global omni_llm
# This will be set by the run_server function
pass
@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
"""Generate text or multimodal content."""
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,
}
]
else:
stage_args = request.stage_args
# Generate using omni_llm
try:
outputs = await omni_llm.generate_async(stage_args)
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:
for out in output.outputs:
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": []})
return GenerateResponse(
outputs=response_outputs,
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))
@app.get("/health")
async def health_check():
"""Health check endpoint."""
return {"status": "healthy", "service": "vllm-omni"}
@app.get("/info")
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
),
"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,
}
for config in omni_llm.stage_configs
],
}
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")
server = uvicorn.Server(config)
await server.serve()
if __name__ == "__main__":
# This is for testing purposes
asyncio.run(run_server(None))
vllm_omni/entrypoints/cli/serve.py
+3 -98
View File
@@ -3,13 +3,8 @@ Omni serve command for vLLM-omni.
"""
import argparse
import asyncio
from typing import List
from vllm_omni.config import DiTConfig, create_ar_stage_config, create_dit_stage_config
from ..omni_llm import AsyncOmniLLM
class OmniServeCommand:
"""Command handler for vLLM-omni serve command."""
@@ -80,103 +75,13 @@ class OmniServeCommand:
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
)
# Start the server
asyncio.run(self._start_server(omni_llm, parsed_args))
pass
def _create_stage_configs(self, args) -> List:
"""Create stage configurations based on arguments."""
stage_configs = []
stage_id = 0
pass
# Add AR stage - use main model if ar_stage not specified
ar_model = args.ar_stage if args.ar_stage else args.model
ar_config = create_ar_stage_config(
stage_id=stage_id,
model_path=ar_model,
input_modalities=["text"],
output_modalities=["text"],
)
stage_configs.append(ar_config)
stage_id += 1
# Add DiT stage if specified
if args.dit_stage:
dit_model = args.dit_stage
elif args.use_diffusers:
# Use a default DiT model if diffusers is enabled
dit_model = "stabilityai/stable-diffusion-2-1"
else:
dit_model = None
if dit_model:
dit_config = DiTConfig(
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=stage_id,
model_path=dit_model,
input_modalities=["text"],
output_modalities=["image"],
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
if "omni" in args.model.lower() or "multimodal" in args.model.lower():
# Assume it's a multimodal model that can handle both AR and DiT
ar_config = create_ar_stage_config(
stage_id=0,
model_path=args.model,
input_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,
)
stage_configs.append(dit_stage_config)
else:
# Default to AR stage
ar_config = create_ar_stage_config(
stage_id=0,
model_path=args.model,
input_modalities=["text"],
output_modalities=["text"],
)
stage_configs.append(ar_config)
return stage_configs
async def _start_server(self, omni_llm: AsyncOmniLLM, args) -> None:
async def _start_server(self, omni_llm, args) -> None:
"""Start the API server."""
try:
# Import here to avoid circular imports
vllm_omni/entrypoints/omni_llm.py
+3 -1
View File
@@ -163,7 +163,9 @@ class OmniStageLLM(LLM):
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, engine_core_output_type=engine_args.engine_output_type
tokenizer=self.llm_engine.tokenizer,
log_stats=self.llm_engine.log_stats,
engine_core_output_type=engine_args.engine_output_type,
)
self.llm_engine.processor = OmniProcessor(
vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
vllm_omni/model_executor/models/__init__.py
+1 -1
View File
@@ -1 +1 @@
from .registry import OmniModelRegistry
from .registry import OmniModelRegistry # noqa: F401
vllm_omni/model_executor/models/qwen2_5_omni.py
+10 -5
View File
@@ -14,6 +14,7 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
@@ -625,7 +626,7 @@ class Qwen2_5OmniForConditionalGeneration(
next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
return next_id.to(dtype=torch.long)
def _init_token2wav_model(self):
def _init_token2wav_model(self, hf_model_folder):
"""Initialize speaker resources if provided; model is constructed in
__init__."""
if self.token2wav is None or self.token2wav_config is None:
@@ -642,7 +643,7 @@ class Qwen2_5OmniForConditionalGeneration(
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
model_path = hf_model_folder
if isinstance(model_path, str) and os.path.isdir(model_path):
spk_pt = os.path.join(model_path, "spk_dict.pt")
if os.path.exists(spk_pt):
@@ -816,10 +817,14 @@ class Qwen2_5OmniForConditionalGeneration(
# Load token2wav weights (if any)
if token2wav_weights and self.token2wav is not None:
self._init_token2wav_model()
hf_model_folder = download_weights_from_hf(
self.vllm_config.model_config.model,
self.vllm_config.load_config.download_dir,
allow_patterns=["*.safetensors", "*.bin", "*.pt"],
)
self._init_token2wav_model(hf_model_folder)
t2w_loaded = self.token2wav.load_weights(
token2wav_weights,
os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"),
token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")
)
t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
loaded_weights.update(t2w_loaded)