update ar_dit_test
This commit is contained in:
hsliu committed
1 parent
ff9c7d4536
commit
81f04d3cf9
32 files changed
+1192
-1908
No files matched your search
@@ -238,7 +238,6 @@ Dockerfile.dev
|
||||
|
||||
# Kubernetes
|
||||
k8s/
|
||||
*.yaml
|
||||
*.yml
|
||||
!configs/*.yaml
|
||||
!configs/*.yml
|
||||
@@ -2,21 +2,22 @@
|
||||
|
||||
## 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.
|
||||
vLLM-omni extends vLLM with multi-stage, multimodal capabilities. It supports both autoregressive (AR) stages and diffusion (DiT) stages, enabling pipelines such as text→image generation.
|
||||
|
||||
## 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
|
||||
1. **vLLM V1 Compatibility** – Reuse LLMEngine/AsyncLLM and executor abstractions.
|
||||
2. **Stage-based Processing** – Each stage encapsulates its own engine and configuration.
|
||||
3. **Executor/Worker Separation** – Executors remain thin; model-specific logic lives in workers and runners.
|
||||
4. **Extensibility** – New modalities can be added by implementing dedicated stage configs, workers, and output processors.
|
||||
|
||||
## 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 --> OutputProcessor (New one for final multimodal output)
|
||||
```
|
||||
Request → OmniLLM/AsyncOmniLLM → StageManager → LLMEngine/Executor → Worker → ModelRunner → OutputProcessor → Response
|
||||
```
|
||||
|
||||
For diffusion stages, the executor hands off to `DiffusionGPUWorker`, which wraps `DiffusionModelRunner` (responsible for diffusers pipeline loading and inference).
|
||||
|
||||
## Core Components (Ordered by Data Flow)
|
||||
|
||||
@@ -27,399 +28,130 @@ 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
|
||||
vllm serve Qwen/Qwen3-0.6B --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
|
||||
### 3. Omni Serve Command (CLI)
|
||||
```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_omni/entrypoints/cli/main.py
|
||||
if "--omni" in sys.argv:
|
||||
omni_args = [arg for arg in sys.argv[1:] if arg not in {"--omni", "serve"}]
|
||||
OmniServeCommand().run(omni_args)
|
||||
else:
|
||||
from vllm.entrypoints.cli.main import main as vllm_main
|
||||
vllm_main()
|
||||
```
|
||||
|
||||
#### vLLM Plugin Integration
|
||||
File: vllm_omni/pyproject.toml
|
||||
### 4. Offline Inference (OmniLLM)
|
||||
```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
|
||||
|
||||
# vllm_omni/entrypoints/omni_llm.py
|
||||
class OmniLLM(LLM):
|
||||
"""Extended LLM supporting multiple engines and stage-based processing"""
|
||||
|
||||
def __init__(self, stage_configs: List[OmniStageConfig]):
|
||||
super().__init__()
|
||||
self.stage_configs = stage_configs
|
||||
self.engine_list = [] # List of AsyncLLM instances for each stage
|
||||
def __init__(..., stage_configs):
|
||||
super().__init__(model=stage_configs[0].model_path if stage_configs else "test-model")
|
||||
self.stage_manager = StageManager(stage_configs)
|
||||
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"""
|
||||
self.stage_manager.initialize_engines()
|
||||
|
||||
# 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
|
||||
def generate(...):
|
||||
current_output = None
|
||||
for idx, (config, args) in enumerate(zip(self.stage_configs, stage_args_list)):
|
||||
engine = self.stage_manager.get_engine(idx)
|
||||
processed = self._process_stage_inputs(config, args, current_output)
|
||||
stage_output = self._execute_stage(engine, processed, stage_config=config)
|
||||
current_output = stage_output
|
||||
config.stage_output = stage_output
|
||||
return self.output_processor.process_output(current_output)
|
||||
```
|
||||
|
||||
### 4. Stage config for setting of different model stages
|
||||
### 5. Stage Manager
|
||||
```python
|
||||
# vllm_omni/core/stage_manager.py
|
||||
class StageManager:
|
||||
def initialize_engines(self):
|
||||
if self._initialized:
|
||||
return
|
||||
for config in self.stage_configs:
|
||||
if config.engine_type == "AR":
|
||||
engine = self._create_ar_engine(config)
|
||||
elif config.engine_type == "DiT":
|
||||
engine = self._create_dit_engine(config)
|
||||
self.engine_list.append(engine)
|
||||
self._initialized = True
|
||||
```
|
||||
|
||||
### 6. Diffusion Execution (DiT Stage)
|
||||
```python
|
||||
# vllm_omni/worker/diffusion_model_runner.py
|
||||
class DiffusionModelRunner:
|
||||
def __init__(self, model_path, pipeline_name=None, device=None, dtype=None):
|
||||
self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype)
|
||||
self._pipeline = self._load_pipeline(model_path, pipeline_name)
|
||||
|
||||
def generate(...):
|
||||
output = self._pipeline(...)
|
||||
return DiffusionRunnerOutput(prompt=prompt, images=getattr(output, "images", []) or [])
|
||||
```
|
||||
|
||||
```python
|
||||
# vllm_omni/worker/gpu_diffusion_worker.py
|
||||
class DiffusionGPUWorker:
|
||||
def __init__(self, model_path, ...):
|
||||
self.model_runner = DiffusionModelRunner(model_path, ...)
|
||||
|
||||
def generate(...):
|
||||
return self.model_runner.generate(...)
|
||||
```
|
||||
|
||||
```python
|
||||
# vllm_omni/executor/diffusers_executor.py
|
||||
class DiffusersPipelineExecutor:
|
||||
def __init__(self, model_path, ...):
|
||||
self.worker = DiffusionGPUWorker(model_path, ...)
|
||||
|
||||
def generate(...):
|
||||
return self.worker.generate(...)
|
||||
```
|
||||
|
||||
This mirrors the standard executor→worker→runner layering in vLLM.
|
||||
|
||||
### 7. Stage Configuration
|
||||
```python
|
||||
# vllm_omni/config/stage_config.py
|
||||
@dataclass
|
||||
class OmniStageConfig:
|
||||
"""Configuration for a processing stage"""
|
||||
stage_id: int
|
||||
engine_type: str # "AR" or "DiT"
|
||||
engine_type: Literal["AR", "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
|
||||
vllm_config: Optional[VllmConfig] = None
|
||||
executor_class: Optional[type[Executor]] = None
|
||||
dit_config: Optional[DiTConfig] = 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 = OmniStageConfig()
|
||||
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 = OmniStageConfig()
|
||||
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
|
||||
- AR stages default to `UniProcExecutor`.
|
||||
- DiT stages set `executor_class` to `DiffusersPipelineExecutor` when `use_diffusers=True`.
|
||||
|
||||
# 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 import AsyncLLM
|
||||
|
||||
|
||||
class AsyncOmniLLM(AsyncLLM):
|
||||
"""Extended AsyncLLM supporting multiple engines and stage-based processing"""
|
||||
|
||||
def __init__(self, stage_configs: List[OmniStageConfig]):
|
||||
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 SchedulerInterface
|
||||
|
||||
|
||||
class OmniDiffusionScheduler(SchedulerInterface):
|
||||
|
||||
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. 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 `OmniDiffusionModelRunner`; 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 OmniDiffusionModelRunner.
|
||||
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
|
||||
from vllm.v1.worker.diffusion_model_runner import OmniDiffusionModelRunner
|
||||
|
||||
class Worker(GPUWorker):
|
||||
def init_device(self) -> None:
|
||||
#those related to device check and init
|
||||
...
|
||||
|
||||
self.model_runner = OmniDiffusionModelRunner(self.vllm_config, ...)
|
||||
```
|
||||
|
||||
### 10. 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`：
|
||||
- OmniDiffusionModelRunner: Use the `pooler_output=[Tensor,...]` to return multi modal tensors
|
||||
- OmniARModelRunner: 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 OmniDiffusionModelRunner(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 OmniARModelRunner(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,...]
|
||||
```
|
||||
|
||||
### 11. 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)
|
||||
```
|
||||
|
||||
|
||||
### 12. OutputProcessor
|
||||
|
||||
For hidden state output, original OutputProcessor already support.
|
||||
|
||||
We only need to add new one for final multimodal output.
|
||||
### 8. Output Processing
|
||||
```python
|
||||
# vllm_omni/engine/output_processor.py
|
||||
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)
|
||||
def process_output(self, engine_core_output):
|
||||
output_type = self._detect_output_type(engine_core_output)
|
||||
handler = self.output_handlers.get(output_type, self._process_pooling_output)
|
||||
return handler(engine_core_output)
|
||||
```
|
||||
|
||||
### 9. CLI Package Configuration
|
||||
```toml
|
||||
# pyproject.toml
|
||||
[project.scripts]
|
||||
vllm = "vllm_omni.entrypoints.cli.main:main"
|
||||
vllm-omni = "vllm_omni.entrypoints.cli.main:main"
|
||||
```
|
||||
|
||||
|
||||
## Future Work
|
||||
- Scheduler adaptations for diffusion batching.
|
||||
- Advanced output processors for audio/video.
|
||||
- Multi-worker execution (ray / multiprocess) for heavy DiT stages.
|
||||
@@ -136,4 +136,4 @@ Hongsheng Liu: [hsliuustc@gmail.com](mailto:hsliuustc@gmail.com)
|
||||
|
||||
[image10]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAXYAAAFiCAYAAADiGfPoAAB7FElEQVR4XuydB1gV19aGc2tuv3/KTUwxsSb2rthFELtSFMUKduzYC1aw94rSRFBBUUBFEVEEsffeYktssbfEWKJ+//4WzvF4wIIVcH/Ps55zZs2ePYXDO2vW7PIBtLS0tLQylT6wdGhpaWlpZWxpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBraWlpZTJpsGtpaWllMmmwa2lpaWUyabBrvbJOnDiBDh06IEuWLPjkk09QpkwZxMTE4N69e5ZF37muXbsmx2dtbY2HDx9artbSyhTSYNd6JY0YMQIffPBBqlaqVCmBfnrS1atXkTdvXhQpUkSDXSvTSoNd66UVGhpqgnixYsWwefNmHDt2DD169DD5K1WqhN9++02i999///2J7R88eJDCd//+/RQ+LtNY3lzG9gQ0zSj3LKUGdvPtUquDx0RLTUZ54zhSk3mdxvlZnotR5mn70dJKizTYtV5Kd+7cQenSpQXe/LQE1ahRo/DHP/5R1q9evRr//ve/5fuyZctkPSFYoEAB8a1bt058BNv//vc/8d29excnT56Eo6Oj6Sbx1Vdf4fDhw6Z9uLm5ib9hw4ZyYzHKtW3bVlIuqckS7KdOncIf/vAH2W7q1KmmOpo2bYpffvkF48aNk+U//elPst4c3tHR0abyxnXYvn27af2NGzfQrl0703pnZ2dUrFhRvk+ePNlULjIy0lTmP//5j1w7nr+W1stKg13rpXTw4EEBEaG4atUqy9USpWfLlk3K9OzZE35+fvLd1tZW1u/cudMEMwKaioqKkuXixYtj3759JuDyCYCAtbKykuXExEQp37FjR1Mdffv2Ffvb3/4my23atDEdi7kswX7mzBnTDejDDz/E2LFj8fXXX5vq5bFw33/+85/leBYvXiz1+Pv7y/r4+HhZDgkJkeWPP/5Y9nHz5k3897//FR/3xzrc3d1N9c6cOVO243FyuXXr1lKGuX8uFy1aVK6hltbLSINd66WUkJAgAGIku2nTJsvVknr57rvvpAyBxuibAP373/8uwDJy8/R98803EvG3atVKfN7e3ujSpYt8b968uUTfBOWWLVvER9BTBti5H0NNmjQRn4uLi8lnrmeBPSAgQMrwBsHl//u//8OFCxfElzt3bvEx/UTdvn0bt27dkjp4bEuWLJH1fHnMbfjymDeCf/7zn3LuhniDYDne6BiVcx9cvnLlitRz9OhR09PNgQMHTNtpaaVFGuxaL6VDhw4JfAivuLg4y9UCPiNinzZtmgDQSEMwHWNnZ4d//etfqFmzpvi2bdsmqZm//vWvcqNwcHAQf2pWvnx52YcBdnt7e9N+27dvL75GjRqZfOZ6Fth37NghZYz0C1/+GqkX4yYVFhZmqisiIkKOkzcr49gMsDPVwmWmiMzTN4ULFxY/I/5Lly7JU4Ll+RnGVI+W1stIg13rpcQI88svvxQAeXh4iI8Aq1evnsCSKQpC+h//+Ad2794t6728vEwgzpo1K/LkyWOK3F1dXSV6zZcvn+S2jdx6586dpT7mws+ePYuff/5Z9k0ZYGeO3ZCR7ngZsDP9QzEdw+WSJUs+FeyEP29qjMC5r9GjR8t6A+w+Pj6yzHM0fyHKfdNPsF++fFnAzqce4xz5yXOk8T2GltbLSINd66VE4BnRMSPvBQsWSPqFIKOPLzr56enpadpm48aNphw4jfBes2aN3AAMX4sWLaRsnz59ZJkpkIsXL4qPTScJxL1798qyAfYGDRqY9vEqYDfqfR7YCeoKFSrIsgHtDRs2yDLBzuPluRrvCJYuXSplmHoxfEzF8Hp9+umnsswUDMWbF1NCQUFBcoPT0noZabBrvbTOnz9vitoJR0bbRqsWw9gZyGgpQrB9//33pnVM0TBqNeBGW758uZQ9fvy4KU3BfRQsWFCgyWWjvncFdvr4Etg4Fp4X00Nc5svT06dPC/D54pU+5tl5/N9++63pPI2Xp0OGDJFlXjeW+fzzz2WZzUR1xK71stJg13plnTt3Dl27dpXWJART9erVpZWIAUOmUwzNmjVLyhByRkqFL0Ppq1y5sqkcRThOmjQJ2bNnl/VsTmjexpwvOeln80ZDvXv3Fh8Bn5quX78uECaYCWge+xdffCHbsKUPxTQKl6tVq2YCe9myZcXHljsUj437582GrXUYsbM5I8vwqcIQmzLmzJlT9sGnGsKb14TXwRCPgWW4LdfzKUZL61Wkwa71RsWUQnrrffo2xDbsQ4cOlTbp5m3bjRYvvBFoab0pabBrab0B8cnAaC3DF6wcR+ezzz6TZUbmT+tApaX1OqTBrqX1hsS27IUKFTLl1Wls/mh0sNLSelPSYNfSeoNip6MffvhB2v3TdEsXrbchDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUwmDXYtLS2tTCYNdi0tLa1MJg12LS0trUymNw72n3/+GQcOHMDhw4fTrZ06dcrysLW0tLQyrN442Ldt24YCBQogd+7c6c9y5ZaJkhMSEiwPW0tLSyvD6o2DfdeuXciZKydGBAzDpAUTMDFsfLqxyeGT4NzaGevWrbM8bK10pNu3b8t0cosXL053xtmRtLTSm94K2HPkzAGfxdOw4XwS1p5Zk25s06X1cGjuoMGeznXmzBm0b99eAP/rr7+mG7ty5QpGjx5tebhaWu9cGuwa7OleBHubNm0s3e9cvNFosGulR2mwa7Cne6VXsP/2228a7FrpUhrsGuzpXhrsWlppkwa7Bnu6lwa7llba9NbBvu5cIjZcSML682tlOelsQgrgvi3TYM8Y0mDX0kqb3irYN15ch4FTPNGqR0ssO7AU/cb1wZDpg5FwKh7L1fKKQ8uw+sRKLNkThVXHYwX6SecSsOLwMsQeXi43BQKZ62NUWa5fdSxWLUdi5dEYKRt/Mg5L90YhUdXJ5dXHWV+k1JtwarXaLhrR+xYj8XR8qmDnC7HAwEDEx8ebnYXWu9TLgD0kJAQ5c+bEZ599hhw5cqBjx46WRZ4QW7i4uLhYup8pDXat9Kp3AvZeo3ti373dGDV7BDoP6giPoV1QuqIVBk31RHXHarBvXBe2tSoL0BdsCkXdRnVQsWoFLNq2AF29uqKWc02UrmAlkHZs6oD6bvVQt3FtKd+oXSNUc6wq++HNoUptW9m+qr0d5q4NwbfZvoVr52aI/zEOmy9vELDzGO/evQtfX198//33+NOf/oStW7danorWO9LLgN1QlSpVTN8J+759++LYsWMIDw/HzZs3xT9//ny5mfMGcP/+fVP550mDXSu96q2DfXiAt3RYKlKyCKxrWCN6/2J0GdwJg6cNkvVNOjTG1qsb0bZ3Gwzz90KP4d0wOXwiVv4QgzmJs1FZbcPIfXrkVFXnVNRr7iTRNyG/4nA0qjtUk2h9xZFl0gmp04AO2HFzK5p3aga/5TNRzaGqKdfPiL1eCye0bNkSefLkwQcffCBGGBw8eBA7d+7E7t27sWfPHuzfv198HILg6NGjOH78OE6ePImffvoJp0+fxtmzZ3H+/HlcuHABly9flgjw+vXruHHjBm7duiUQeJrxKeFpdufOnWcab0hPs3v37qVqv//++1ONYHuaPXjw4Kn28OHDZ9qr6Ny5cy8N9kqVKsm1iIqKQpcuXZCUlIQKFSrI37B3795YsmQJpk+fLh2N6E/LsWqwa6VXvXWwS8Q+qid2/7YD265vxrqfE9F9mAcmhU9A/Ik4tOjmho2X1sFzUn94zRiCVt1bwCdqGrZe24SI7QtRupIV7JvYo55rPSzbvwQzlk5HdadqaNaxCVafjMP8DfMkOh8ROAwj1E2kx4ju2Hp9k9pnD8xc6gOXNg0lv2+AvX6LeqhXrx4++ugjE9hz5cqFfv36oUePHujWrRs8PDzQuXNndOrUSTrKtG3bFq1bt0arVq3QokULuLq6olmzZmjSpAkaNWqEhg0bwtnZWep1dHRE3bp1YW9vL+bg4CBGP83JyUmMZevXry/WoEEDMdZDY4qAxrpp3I9hTZs2FeP+mzdvLsbjcXNzE+Px0XisNB43jaCk8Vxo7dq1g7u7uxjPsUOHDmJMYfC8abwGNAKya9euYrw2NF6n7t27yzWj9ezZU6xXr15ihGifPn0kYjasf//+Yp6enmIDBgwQGzRoEAYOHCifgwcPlv3w+ryMCHbezLhvXq/x48fLcBKUv78/6tSpI995U65Ro4b5ps+VBrtWetU7AXtnFaEb+XLmwbt5d5Xomi9VrWtUknw44cwIu8+Y3ug/sR96j+mFoT6DUc2xmmw7ft5YTI+YqsDdDbtubYd7v7YCbkb6i7YuUJ+tVYQfLPWs+XEVKtpVQHB8EBq0cn4C7EzFbNq0SaJuAurvf/87smbNikuXLlmeyjuXEf0yQmYEbUTZBBejUkbwRqRP6PBJgcZekr/88oukHmh8iuDTxLVr18SuXr0qTxh80uB5X7x4UYxPHwQeB3KjMXLmkwlTIzReMw6gRuOTy48//ihPMSdOnJAnGhrTHoyOGRHTjhw5Ik89hw4dEuNTEAeJo/GpaN++fdi7d688JdH4xBQbGys3q5eRAXYvLy9MnDgRGzZskCcxasyYMXKz5fXi+Wqwa2UWvV2wX1iHKQsnYeSs4abWMPwcEzwSs+ICJHof5jcUFezKq4i9n9wIWLine truncated
|
||||
|
||||
[image11]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABM8AAARlCAYAAACukOBGAACAAElEQVR4Xuy9B7hV13WuHSd5ch3Zia17rcQl13+K88dyXOLYlmPLRb13CVUkUEMIEB3Rkeiig+iiI3rvvffeO4gmeu+dM+56xz4TbRZ7IxDt7M33Ps98zll97rXGnGvOb4055l+YEEIIIYQQQgghhBAiJX8RXyGEEEIIIYQQQgghhEgg8UwIIYQQQgghhBBCiDRIPBNCCCGEEEIIIYQQIg0Sz4QQQgghhBBCCCGESIPEMyGEEEIIIYQQQggh0iDxTAghhBBCCCGEEEKINEg8E0IIIYQQQgghhBAiDRLPhBBCCCGEEEIIIYRIg8QzIYQQQgghhBBCCCHSIPFMCCGEEEIIIYQQQog0SDwTQgghhBBCCCGEECINEs+EuAhnz5jlKCnF09k0Kb49/B9fH98/1bazSfaXYhvr025PPm98e3xd8nJ826VcK36eVNeN/e/nih+Xajl+vaT1550jvu+lbIv/TfV/fDn5nPHz5yYhhBBCCCFE9iHxTIh05JgdP5Bjx3YrKcXSntTpeG5KXk61/lzafZFtsXPE07n1KfJz7pwp8nvB9ZL2uWBb8j7ptuWm+DWT973geqmOS943zfW+7Den+j++HD9Hym25/8eXk895wfmj8506mhOvRYQQQgghhBBZgMQzIdKBeLY/x45sz7GjO5SUvkhH0qT49vB/fH18/1TbLnbclx2bvC2+Pb4ueTm+Lb5Pqm2p9onvG///Upfj10ten25b/P/4cvzcqbalukb8nBecf7vEMyGEEEIIIbIViWdCXIQT+y3ROd6ppKSklD5RT0g8E0IIIYQQIjuReCZEOnIQz3ISHigpOstKSkpKIUk8E0IIIYQQInuReCZEOsKwTYlnSkpKX5IkngkhhBBCCJG9SDwTIh3yPFNSUrrEJPFMCCGEEEKI7EXimRDpkOeZkpLSJSaJZ0IIIYQQQmQvEs+ESIfEMyUlpUtMEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBPiIpzYbxLPbmA6vjt6CEej57Dni+dgh83OHkgsu2CxL9rnpPlflk9wzIlon4Nf7HNyb2IfO56bonMe22UXXC+euD7nuZR9SX79KK/k8cuOYd/TkX3lHJKNZUOSeCaEEEIIIUT2IvFMiHRowoAbmhC8tq3bby0afWKfLd3mQhPrenToZ2MGT3ax7PRBs3lTllmJImVsxbz1dvpQ4pjypSrb8H7jfDv7LZu91t59s5gVzk1Vy9ew3RtQuOycoIZQhuCFsBYEtu1r91vfbkNs/+fsYC7K2ZGEKOYiGeuOJY4NQt6GKK9d2va0PZuOfnGu4wkB0MXAIOBFl58zabGNHDDBBTpSOD/58GXyd+ILYZB74PvkCm5nDnyxHL9/Stc3STwTQgghhBAie5F4JkQ6NGzzhiaEorWLt9hD9z9mcycvdVGJde++Ucwa1m7m4hMi1IAew+373/uBNW/Y1h/bkN6jfblRneYuPuHZNbDHCLvvrgds5viFtmDaUlsya5Xt2XjE5kxaZD079bcB3YfZwa2nXOBCdOvZaUC03zL7fNVuGzFgvK1bssUmjphhndv0sCmjZrvnG/v36NjPBvYc4fk9sSfHBa/5U5ZZwfxv2cblO23RjOXWu8tA69NlkO3//IQd3n7Gl7u172N7N+2zSmWq2aMPPm57Nh2wxTNXW6fW3WzCsOkuni2ZtSbK93Dr3qFvlI89djr6vRuWbY/y0N1mTljowtyaRZs9T3OnLPU8xe+h0vVLEs+EEEIIIYTIXiSeCXERNGzzxiWEMjzO8j39ki2YtsLFMtaVKVbeWjRs66IYAtKwvmPtyUeftVJFy9mhHcetecM2lu+pF619y09dcEPQGtpnjL30XH4XoXatP2CHo/3WR+f+nzvudE+0P995t/XpOtjGD5saXe9Fe+O1QvbgfY/YzAkLrGzxCtamWQe741f/YyXeLWP3/Pk+Wz5njbVq3N5KFStnBV9509o07eBeYORp0fSV5zzh7r/nQXvztXeiY+73/HRq3d3ee6ekVS73gQtolcpWs8ceftKFOfLfpF5Le+bxfDZ93Hx75YUC9nz025945Gmr8n4N2711X/S3ur1fspIVi84xacxUK/1eeav9QT17KV9+mzNxkZ2WgHbDksQzIYQQQgghsheJZ0JcBIlnNy5dqng2uNcoK1mkjL1fqrIN6D7CalerbzUq17HWTTu4eMZ+YwZNstv/4z8t/wsFrMDLr1vLxm1s7aLP7bmnXrAd6w5Yp1bdrH7NJlauRAUfErp1zT4r9EYRmzZmrpUrWTHa/xN7r3ApP1fZaJ9P2/exJx992lo1aW8fVKhlTz/+nB3YctLzE8SzZXPWuaj12dKt1rfrYKsV5avtx52sepS3kQNG277NB21g9+HWqM7HdmLP2SjvQ21Qr5H2wrMvW99Ph1iRQiVsanT9GePmW8WyVaPtw/28eNMd2XPUOrfuYX+68y4b2meYPfLA4+5pR/7i91Hp+iSJZ0IIIYQQQmQvEs+ESIeGbd7QhFC2bulWe+HZV1zock4kxLMOLbuee0yDeo6wmlU/cq8uvLVqVatnrZu0TwzjzPU8Q2B7u0DhRKyxXH1j3eIt7nW2c/2h6HyfutcXItzEETNt+7oDVrRQiYR4VqKifdK8k4tfDKdkqGW39n3tuadfsJpVPrKPG7T0oZOHtp2xnMNfiGcMu+TvphU7rVfHAVbnw4Z2cNux6Nge7s3W9ZNePuyzWb1WNnviYnv5+VetWf1W9uC9D/tQ0MrlPrSF01fYhOHTo99X14eSFn+39Lk4bd2jPNz9p/tcCKxfs7FNGjnTTu298D4qXZ8k8UwIIYQQQojsReKZEOmQeHZDE+LZ+mXb7Q+/v8uFq95detnMCXOtyFvF7eV8r1qvzr1sxMARLnzh/TV2yBS77Tu32bC+Y6xFo7ZWv1ZT9wTDGwvx7Df//VsX2Lq1722Deg6xuZOXWIGX37Bta/db8wZtrHHdFja490h7q0Bhq/NBA3vmiXw2bew8K/p2cR8KWrFMNc8XotqwfmPdY6xhnWY+bJLjT+xNzARKzLO3Cha2hdNX+vDPz5ZstS5te9iHlWpbuxado3O19mM/ad7Z+n06xAq88oa1b9HFr9e/21B75IHHXIwrWbSsTygwatBEK1uivK2cv87PhyhYqlhZ69Wlr4tzvToPsBLvlrbJo2b5PYvfR6XrkySeCSGEEEIIkb1IPBMiHRLPbmg6vjvH9m85YZ+262N1PmxgtarWs8G9R3nA/no1Gvsyccfw2po1YaHt3njExadd6w/ZwmkrfBbOk/uYtTPHA+03rdfSaler555pzeq3ttULNrlXF9fAW2zBtOW2d9Mxa1i7oRV45U2PZcZxU0fPskUzVvrwSfJFfLItq/b49Zi4oEGtprZpxS47tS8xacC2Nfv8vMzUOX7YNNuz8agtn7POxbrN0X51qzf0vByIrrt55S6Pp8YQT2boRFQbPWiSX2/K6Nm2dfVej82GMMaEALMmLvJhqUP6jLacI4nZOvk9vToNiOz0rN+z+H1Uuj5J4pkQQgghhBDZi8QzIdIh8eyGJ4ZJMlTTTuemI5aYZTMsn0gMy8yJ0ond5sMZmTGT4P2ITTw70sm9SceQTpqd2pc4F9dgf87F7JYli5axd94o4gIY52afMBkAeWKZ83EdO5VILIdr+frDudujv8d3J2K1cS7PB8ecTKz35ZOJ7f47o//5jZ6fQ4lznGKfQ4nze5y33PtwlGXyx/KxxO+I3z+l65ckngkhhBBCCJG9SDwTIh0Sz266hACFCIewpiGQSpeTJJ4JIYQQQgiRvUg8EyIdEs9uyoSAJi8upctNEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBMiHRLPlJSULjFJPBNCCCGEECJ7kXgmRDoknikpKV1ikngmhBBCCCFE9iLxTIh0SDxTUlK6xCTxTAghhBBCiOxF4pkQ6ZB4pqSkdIlJ4pkQQgghhBDZi8QzIdIh8UxJKU8kZj89nsdnQJV4JoQQQgghRPYi8UyIdCCe7Ys6xdsTHWMlJaXrn45G6cRus5N77YJteSpF9cSpIxLPhBBCCCGEyEYkngmRhpyoH4wnCR1iJaXz0lG7cDms+7JtydtTLce3pTrvxY5LtRzfFj82+RrxY+P7xrelOzbVvun+j58nadvZE2Zd2/ewMsUr2IlDZ+x0/Nj4NeLnSZe/VMelOybVtvj/UTpzQuKZEEIIIYQQ2YjEMyGEEHmacePHWdu2bS0n52x8kxBCCCGEEEJccySeCSGEEEIIIYQQQgiRBolnQgghhBBCCCGEEEKkQeKZEEIIIYQQQgghhBBpkHgmhBAiTzNixAj76KOP7OxZxTwTQgghhBBCXH8kngkhhMjTFC9e3G699VY7ffp0fJMQQgghhBBCXHMkngkhhMjTVKxY0f7t3/7NTp06Fd8khBBCCCGEENcciWdCCCHyNJUqVbJ//dd/lXgmhBBCCCGEuCFIPBNCCJGnqVy5sv3Lv/yLxDMhhBBCCCHEDUHimRBCiDwNwzZ/+MMfSjwTQgghhBBC3BAkngkhhMjT4Hkm8UwIIYQQQghxo5B4JoQQIk+za9cuW7NmjeXk5MQ3CSGEEEIIIcQ1R+KZEEIIIYQQQgghhBBpkHgmhBBCCCGEEEIIIUQaJJ4JIYQQQgghhBBCCJEGiWdCCCHyNK1bt7Z77rlHEwYIIYQQQgghbggSz4QQQuRpqlSpYn//939vJ0+ejG8SQgghhBBCiGuOxDMhhBB5mpo1a9p3v/tdiWdCCCGEEEKIG4LEMyGEEHmaWrVq2T/+4z9KPBNCCCGEEELcECSeCSGEyNMgnv3DP/yDxDMhhBBCCCHEDUHimRBCiDxNtWrV7JZbbpF4JoQQQgghhLghSLine truncated
|
||||
[image11]: <data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABM8AAARlCAYAAACukOBGAACAAElEQVR4Xuy9B7hV13WuHSd5ch3Zia17rcQl13+K88dyXOLYlmPLRb13CVUkUEMIEB3Rkeiig+iiI3rvvffeO4gmeu+dM+56xz4TbRZ7IxDt7M33Ps98zll97rXGnGvOb4055l+YEEIIIYQQQgghhBAiJX8RXyGEEEIIIYQQQgghhEgg8UwIIYQQQgghhBBCiDRIPBNCCCGEEEIIIYQQIg0Sz4QQQgghhBBCCCGESIPEMyGEEEIIIYQQQggh0iDxTAghhBBCCCGEEEKINEg8E0IIIYQQQgghhBAiDRLPhBBCCCGEEEIIIYRIg8QzIYQQQgghhBBCCCHSIPFMCCGEEEIIIYQQQog0SDwTQgghhBBCCCGEECINEs+EuAhnz5jlKCnF09k0Kb49/B9fH98/1bazSfaXYhvr025PPm98e3xd8nJ826VcK36eVNeN/e/nih+Xajl+vaT1550jvu+lbIv/TfV/fDn5nPHz5yYhhBBCCCFE9iHxTIh05JgdP5Bjx3YrKcXSntTpeG5KXk61/lzafZFtsXPE07n1KfJz7pwp8nvB9ZL2uWBb8j7ptuWm+DWT973geqmOS943zfW+7Den+j++HD9Hym25/8eXk895wfmj8506mhOvRYQQQgghhBBZgMQzIdKBeLY/x45sz7GjO5SUvkhH0qT49vB/fH18/1TbLnbclx2bvC2+Pb4ueTm+Lb5Pqm2p9onvG///Upfj10ten25b/P/4cvzcqbalukb8nBecf7vEMyGEEEIIIbIViWdCXIQT+y3ROd6ppKSklD5RT0g8E0IIIYQQIjuReCZEOnIQz3ISHigpOstKSkpKIUk8E0IIIYQQInuReCZEOsKwTYlnSkpKX5IkngkhhBBCCJG9SDwTIh3yPFNSUrrEJPFMCCGEEEKI7EXimRDpkOeZkpLSJSaJZ0IIIYQQQmQvEs+ESIfEMyUlpUtMEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBPiIpzYbxLPbmA6vjt6CEej57Dni+dgh83OHkgsu2CxL9rnpPlflk9wzIlon4Nf7HNyb2IfO56bonMe22UXXC+euD7nuZR9SX79KK/k8cuOYd/TkX3lHJKNZUOSeCaEEEIIIUT2IvFMiHRowoAbmhC8tq3bby0afWKfLd3mQhPrenToZ2MGT3ax7PRBs3lTllmJImVsxbz1dvpQ4pjypSrb8H7jfDv7LZu91t59s5gVzk1Vy9ew3RtQuOycoIZQhuCFsBYEtu1r91vfbkNs/+fsYC7K2ZGEKOYiGeuOJY4NQt6GKK9d2va0PZuOfnGu4wkB0MXAIOBFl58zabGNHDDBBTpSOD/58GXyd+ILYZB74PvkCm5nDnyxHL9/Stc3STwTQgghhBAie5F4JkQ6NGzzhiaEorWLt9hD9z9mcycvdVGJde++Ucwa1m7m4hMi1IAew+373/uBNW/Y1h/bkN6jfblRneYuPuHZNbDHCLvvrgds5viFtmDaUlsya5Xt2XjE5kxaZD079bcB3YfZwa2nXOBCdOvZaUC03zL7fNVuGzFgvK1bssUmjphhndv0sCmjZrvnG/v36NjPBvYc4fk9sSfHBa/5U5ZZwfxv2cblO23RjOXWu8tA69NlkO3//IQd3n7Gl7u172N7N+2zSmWq2aMPPm57Nh2wxTNXW6fW3WzCsOkuni2ZtSbK93Dr3qFvlI89djr6vRuWbY/y0N1mTljowtyaRZs9T3OnLPU8xe+h0vVLEs+EEEIIIYTIXiSeCXERNGzzxiWEMjzO8j39ki2YtsLFMtaVKVbeWjRs66IYAtKwvmPtyUeftVJFy9mhHcetecM2lu+pF619y09dcEPQGtpnjL30XH4XoXatP2CHo/3WR+f+nzvudE+0P995t/XpOtjGD5saXe9Fe+O1QvbgfY/YzAkLrGzxCtamWQe741f/YyXeLWP3/Pk+Wz5njbVq3N5KFStnBV9509o07eBeYORp0fSV5zzh7r/nQXvztXeiY+73/HRq3d3ee6ekVS73gQtolcpWs8ceftKFOfLfpF5Le+bxfDZ93Hx75YUC9nz025945Gmr8n4N2711X/S3ur1fspIVi84xacxUK/1eeav9QT17KV9+mzNxkZ2WgHbDksQzIYQQQgghsheJZ0JcBIlnNy5dqng2uNcoK1mkjL1fqrIN6D7CalerbzUq17HWTTu4eMZ+YwZNstv/4z8t/wsFrMDLr1vLxm1s7aLP7bmnXrAd6w5Yp1bdrH7NJlauRAUfErp1zT4r9EYRmzZmrpUrWTHa/xN7r3ApP1fZaJ9P2/exJx992lo1aW8fVKhlTz/+nB3YctLzE8SzZXPWuaj12dKt1rfrYKsV5avtx52sepS3kQNG277NB21g9+HWqM7HdmLP2SjvQ21Qr5H2wrMvW99Ph1iRQiVsanT9GePmW8WyVaPtw/28eNMd2XPUOrfuYX+68y4b2meYPfLA4+5pR/7i91Hp+iSJZ0IIIYQQQmQvEs+ESIeGbd7QhFC2bulWe+HZV1zock4kxLMOLbuee0yDeo6wmlU/cq8uvLVqVatnrZu0TwzjzPU8Q2B7u0DhRKyxXH1j3eIt7nW2c/2h6HyfutcXItzEETNt+7oDVrRQiYR4VqKifdK8k4tfDKdkqGW39n3tuadfsJpVPrKPG7T0oZOHtp2xnMNfiGcMu+TvphU7rVfHAVbnw4Z2cNux6Nge7s3W9ZNePuyzWb1WNnviYnv5+VetWf1W9uC9D/tQ0MrlPrSF01fYhOHTo99X14eSFn+39Lk4bd2jPNz9p/tcCKxfs7FNGjnTTu298D4qXZ8k8UwIIYQQQojsReKZEOmQeHZDE+LZ+mXb7Q+/v8uFq95detnMCXOtyFvF7eV8r1qvzr1sxMARLnzh/TV2yBS77Tu32bC+Y6xFo7ZWv1ZT9wTDGwvx7Df//VsX2Lq1722Deg6xuZOXWIGX37Bta/db8wZtrHHdFja490h7q0Bhq/NBA3vmiXw2bew8K/p2cR8KWrFMNc8XotqwfmPdY6xhnWY+bJLjT+xNzARKzLO3Cha2hdNX+vDPz5ZstS5te9iHlWpbuxado3O19mM/ad7Z+n06xAq88oa1b9HFr9e/21B75IHHXIwrWbSsTygwatBEK1uivK2cv87PhyhYqlhZ69Wlr4tzvToPsBLvlrbJo2b5PYvfR6XrkySeCSGEEEIIkb1IPBMiHRLPbmg6vjvH9m85YZ+262N1PmxgtarWs8G9R3nA/no1Gvsyccfw2po1YaHt3njExadd6w/ZwmkrfBbOk/uYtTPHA+03rdfSaler555pzeq3ttULNrlXF9fAW2zBtOW2d9Mxa1i7oRV45U2PZcZxU0fPskUzVvrwSfJFfLItq/b49Zi4oEGtprZpxS47tS8xacC2Nfv8vMzUOX7YNNuz8agtn7POxbrN0X51qzf0vByIrrt55S6Pp8YQT2boRFQbPWiSX2/K6Nm2dfVej82GMMaEALMmLvJhqUP6jLacI4nZOvk9vToNiOz0rN+z+H1Uuj5J4pkQQgghhBDZi8QzIdIh8eyGJ4ZJMlTTTuemI5aYZTMsn0gMy8yJ0ond5sMZmTGT4P2ITTw70sm9SceQTpqd2pc4F9dgf87F7JYli5axd94o4gIY52afMBkAeWKZ83EdO5VILIdr+frDudujv8d3J2K1cS7PB8ecTKz35ZOJ7f47o//5jZ6fQ4lznGKfQ4nze5y33PtwlGXyx/KxxO+I3z+l65ckngkhhBBCCJG9SDwTIh0Sz266hACFCIewpiGQSpeTJJ4JIYQQQgiRvUg8EyIdEs9uyoSAJi8upctNEs+EEEIIIYTIXiSeCZEOiWdKSkqXmCSeCSGEEEIIkb1IPBMiHRLPlJSULjFJPBNCCCGEECJ7kXgmRDoknikpKV1ikngmhBBCCCFE9iLxTIh0SDxTUlK6xCTxTAghhBBCiOxF4pkQ6ZB4pqSkdIlJ4pkQQgghhBDZi8QzIdIh8UxJKU8kZj89nsdnQJV4JoQQQgghRPYi8UyIdCCe7Ys6xdsTHWMlJaXrn45G6cRus5N77YJteSpF9cSpIxLPhBBCCCGEyEYkngmRhpyoH4wnCR1iJaXz0lG7cDms+7JtydtTLce3pTrvxY5LtRzfFj82+RrxY+P7xrelOzbVvun+j58nadvZE2Zd2/ewMsUr2IlDZ+x0/Nj4NeLnSZe/VMelOybVtvj/UTpzQuKZEEIIIYQQ2YjEMyGEEHmacePHWdu2bS0n52x8kxBCCCGEEEJccySeCSGEEEIIIYQQQgiRBolnQgghhBBCCCGEEEKkQeKZEEIIIYQQQgghhBBpkHgmhBAiTzNixAj76KOP7OxZxTwTQgghhBBCXH8kngkhhMjTFC9e3G699VY7ffp0fJMQQgghhBBCXHMkngkhhMjTVKxY0f7t3/7NTp06Fd8khBBCCCGEENcciWdCCCHyNJUqVbJ//dd/lXgmhBBCCCGEuCFIPBNCCJGnqVy5sv3Lv/yLxDMhhBBCCCHEDUHimRBCiDwNwzZ/+MMfSjwTQgghhBBC3BAkngkhhMjT4Hkm8UwIIYQQQghxo5B4JoQQIk+za9cuW7NmjeXk5MQ3CSGEEEIIIcQ1R+KZEEIIIYQQQgghhBBpkHgmhBBCCCGEEEIIIUQaJJ4JIYQQQgghhBBCCJEGiWdCCCHyNK1bt7Z77rlHEwYIIYQQQgghbggSz4QQQuRpqlSpYn//939vJ0+ejG8SQgghhBBCiGuOxDMhhBB5mpo1a9p3v/tdiWdCCCGEEEKIG4LEMyGEEHmaWrVq2T/+4z9KPBNCCCGEEELcECSeCSGEyNMgnv3DP/yDxDMhhBBCCCHEDUHimRBCiDxNtWrV7JZbbpF4JoQQQgghhLghSLine truncated
|
||||
@@ -12,13 +12,12 @@ vllm_omni/
|
||||
│ └── cache_config.py # Caching configurations
|
||||
├── core/ # Core processing components
|
||||
│ ├── __init__.py
|
||||
│ ├── omni_llm.py # OmniLLM and AsyncOmniLLM
|
||||
│ ├── stage_manager.py # Multi-stage orchestration
|
||||
│ ├── dit_cache_manager.py # DiT caching system
|
||||
│ └── sched/ # Schedulers
|
||||
│ ├── __init__.py
|
||||
│ ├── scheduler.py # Base scheduler interface
|
||||
│ └── diffusion_scheduler.py # DiT scheduler
|
||||
|
||||
├── engine/ # Engine components
|
||||
│ ├── __init__.py
|
||||
│ ├── processor.py # Output processing
|
||||
@@ -29,14 +28,14 @@ vllm_omni/
|
||||
│ └── diffusers_executor.py # Diffusers pipeline executor
|
||||
├── model_executor/ # Model execution components
|
||||
│ ├── __init__.py
|
||||
│ ├── ar_model_runner.py # OmniARModelRunner
|
||||
│ └── dit_model_runner.py # OmniDiffusionModelRunner
|
||||
├── worker/ # Worker implementations
|
||||
│ ├── __init__.py
|
||||
│ └── omni_worker.py # Extended worker for DiT
|
||||
│ ├── 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
|
||||
@@ -87,24 +86,29 @@ class OmniStageConfig:
|
||||
input_modalities: List[str]
|
||||
output_modalities: List[str]
|
||||
vllm_config: Optional[VllmConfig] = None
|
||||
executor_class: type[Executor] = MultiprocExecutor
|
||||
dit_config: Optional[DiTConfig] = None
|
||||
cache_config: Optional[DiTCacheConfig] = None
|
||||
executor_class: Optional[type[Executor]] = None
|
||||
stage_output: Optional[Any] = None
|
||||
```
|
||||
|
||||
### 3.2 DiT Configuration
|
||||
|
||||
```python
|
||||
# vllm_omni/config/dit_config.py
|
||||
# vllm_omni/config/stage_config.py
|
||||
@dataclass
|
||||
class DiTConfig:
|
||||
model_type: str
|
||||
scheduler_type: str
|
||||
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
|
||||
guidance_scale: float = 7.5
|
||||
use_diffusers: bool = False
|
||||
diffusers_pipeline: Optional[str] = None
|
||||
height: int = 512
|
||||
width: int = 512
|
||||
```
|
||||
|
||||
### 3.3 Cache Configuration
|
||||
@@ -179,34 +183,9 @@ class AsyncOmniLLM(AsyncLLM):
|
||||
"""Async generation interface"""
|
||||
# Similar to OmniLLM but with async/await patterns
|
||||
pass
|
||||
```
|
||||
``
|
||||
|
||||
### 4.3 DiT Scheduler Implementation
|
||||
|
||||
```python
|
||||
# vllm_omni/core/sched/diffusion_scheduler.py
|
||||
class OmniDiffusionScheduler(SchedulerInterface):
|
||||
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:
|
||||
super().__init__(vllm_config, kv_cache_config, structured_output_manager,
|
||||
mm_registry, include_finished_set, log_stats)
|
||||
self.dit_cache_manager = DiTCacheManager(dit_cache_config)
|
||||
|
||||
def schedule(self) -> SchedulerOutput:
|
||||
"""Schedule DiT requests with caching optimization"""
|
||||
# Implement DiT-specific scheduling logic
|
||||
pass
|
||||
```
|
||||
|
||||
### 4.4 Model Runners
|
||||
### 4.3 Model Runners
|
||||
|
||||
```python
|
||||
# vllm_omni/model_executor/dit_model_runner.py
|
||||
|
||||
@@ -1,293 +0,0 @@
|
||||
# vLLM-omni Implementation Summary
|
||||
|
||||
## Overview
|
||||
|
||||
This document provides a comprehensive summary of the vLLM-omni implementation plan and current progress. The implementation follows a structured approach with PRD, architecture design, test design, and code implementation phases.
|
||||
|
||||
## Implementation Status
|
||||
|
||||
### ✅ Completed Phases
|
||||
|
||||
#### 1. Product Requirements Document (PRD)
|
||||
- **File**: `docs/PRD.md`
|
||||
- **Status**: Complete
|
||||
- **Contents**:
|
||||
- Product vision and target users
|
||||
- Functional and non-functional requirements
|
||||
- Technical architecture overview
|
||||
- Implementation phases and success criteria
|
||||
- Risk assessment and mitigation strategies
|
||||
|
||||
#### 2. Architecture Design
|
||||
- **File**: `docs/architecture/implementation_architecture.md`
|
||||
- **Status**: Complete
|
||||
- **Contents**:
|
||||
- Detailed package structure
|
||||
- Core module dependencies
|
||||
- Configuration system design
|
||||
- Implementation details for all components
|
||||
- Memory management and error handling strategies
|
||||
|
||||
#### 3. Test Design
|
||||
- **File**: `docs/testing/test_design.md`
|
||||
- **Status**: Complete
|
||||
- **Contents**:
|
||||
- Comprehensive testing strategy
|
||||
- Unit, integration, and E2E test specifications
|
||||
- Performance and compatibility testing
|
||||
- Test configuration and execution guidelines
|
||||
|
||||
#### 4. Package Setup
|
||||
- **Status**: Complete
|
||||
- **Components**:
|
||||
- Updated `pyproject.toml` with vLLM integration
|
||||
- Package structure and dependencies
|
||||
- CLI entry point configuration
|
||||
- vLLM plugin system setup
|
||||
|
||||
#### 5. Core Modules
|
||||
- **Status**: Complete
|
||||
- **Components**:
|
||||
- `vllm_omni/config/`: Configuration management
|
||||
- `vllm_omni/core/omni_llm.py`: OmniLLM and AsyncOmniLLM classes
|
||||
- `vllm_omni/core/stage_manager.py`: Multi-stage orchestration
|
||||
- `vllm_omni/core/dit_cache_manager.py`: DiT caching system
|
||||
- `vllm_omni/core/sched/diffusion_scheduler.py`: DiT scheduler
|
||||
- `vllm_omni/engine/output_processor.py`: Multimodal output processing
|
||||
|
||||
### 🚧 In Progress / Pending
|
||||
|
||||
#### 6. Scheduler and Executor Components
|
||||
- **Status**: Partially Complete
|
||||
- **Completed**:
|
||||
- DiT scheduler implementation
|
||||
- DiT cache manager
|
||||
- **Pending**:
|
||||
- Executor implementations
|
||||
- Integration with vLLM's executor system
|
||||
|
||||
#### 7. Model Runners
|
||||
- **Status**: Pending
|
||||
- **Components**:
|
||||
- `OmniDiffusionModelRunner`: DiT model execution
|
||||
- `OmniARModelRunner`: AR model execution with hidden states
|
||||
- Integration with vLLM's model runner system
|
||||
|
||||
#### 8. Output Processing
|
||||
- **Status**: Partially Complete
|
||||
- **Completed**:
|
||||
- Basic multimodal output processor
|
||||
- **Pending**:
|
||||
- RequestState extensions
|
||||
- Advanced output handling
|
||||
|
||||
#### 9. CLI Integration
|
||||
- **Status**: Complete
|
||||
- **Components**:
|
||||
- CLI entry point with `--omni` flag support
|
||||
- OmniServeCommand implementation
|
||||
- API server for serving
|
||||
- vLLM plugin system
|
||||
|
||||
#### 10. Testing and Validation
|
||||
- **Status**: Partially Complete
|
||||
- **Completed**:
|
||||
- Basic test structure
|
||||
- Configuration tests
|
||||
- **Pending**:
|
||||
- Integration tests
|
||||
- E2E tests
|
||||
- Performance benchmarks
|
||||
|
||||
## Key Implementation Details
|
||||
|
||||
### Package Structure
|
||||
```
|
||||
vllm_omni/
|
||||
├── config/ # Configuration management
|
||||
│ ├── stage_config.py # OmniStageConfig, DiTConfig, etc.
|
||||
│ └── __init__.py
|
||||
├── core/ # Core processing components
|
||||
│ ├── omni_llm.py # OmniLLM and AsyncOmniLLM
|
||||
│ ├── stage_manager.py # Multi-stage orchestration
|
||||
│ ├── dit_cache_manager.py # DiT caching system
|
||||
│ └── sched/ # Schedulers
|
||||
│ └── diffusion_scheduler.py
|
||||
├── engine/ # Engine components
|
||||
│ └── output_processor.py # Multimodal output handling
|
||||
├── entrypoints/ # Entry points and CLI
|
||||
│ ├── cli/main.py # CLI main entry point
|
||||
│ ├── omni.py # OmniServeCommand
|
||||
│ └── api_server.py # API server
|
||||
├── plugin.py # vLLM plugin system
|
||||
└── __init__.py # Main package exports
|
||||
```
|
||||
|
||||
### Core Classes
|
||||
|
||||
#### OmniLLM
|
||||
- **Purpose**: Offline multi-stage processing
|
||||
- **Features**:
|
||||
- Sequential stage processing
|
||||
- AR and DiT stage support
|
||||
- Input/output processing between stages
|
||||
- Integration with vLLM's LLMEngine
|
||||
|
||||
#### AsyncOmniLLM
|
||||
- **Purpose**: Online multi-stage processing
|
||||
- **Features**:
|
||||
- Asynchronous stage processing
|
||||
- Real-time request handling
|
||||
- Integration with vLLM's AsyncLLM
|
||||
|
||||
#### StageManager
|
||||
- **Purpose**: Orchestrate multiple stage engines
|
||||
- **Features**:
|
||||
- Engine creation and management
|
||||
- Stage configuration handling
|
||||
- Resource cleanup
|
||||
|
||||
#### DiTCacheManager
|
||||
- **Purpose**: Optimize DiT inference with caching
|
||||
- **Features**:
|
||||
- Tensor caching with multiple strategies (FIFO, LRU, LFU)
|
||||
- Memory management
|
||||
- Cache statistics and monitoring
|
||||
|
||||
### Configuration System
|
||||
|
||||
#### OmniStageConfig
|
||||
- **Purpose**: Configure individual processing stages
|
||||
- **Features**:
|
||||
- Support for AR and DiT engine types
|
||||
- Modality specification
|
||||
- Executor class selection
|
||||
- DiT-specific configuration
|
||||
|
||||
#### DiTConfig
|
||||
- **Purpose**: Configure DiT-specific parameters
|
||||
- **Features**:
|
||||
- Model type and scheduler configuration
|
||||
- Inference parameters (steps, guidance scale)
|
||||
- Diffusers integration support
|
||||
|
||||
### CLI Integration
|
||||
|
||||
#### Command Line Interface
|
||||
```bash
|
||||
# Basic usage
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8000
|
||||
|
||||
# Multi-stage configuration
|
||||
vllm serve model --omni --ar-stage text-model --dit-stage image-model
|
||||
|
||||
# DiT-specific options
|
||||
vllm serve model --omni --dit-steps 50 --dit-guidance-scale 7.5
|
||||
```
|
||||
|
||||
#### API Server
|
||||
- **Endpoints**:
|
||||
- `POST /generate`: Generate text or multimodal content
|
||||
- `GET /health`: Health check
|
||||
- `GET /info`: Service information
|
||||
- **Features**:
|
||||
- FastAPI-based REST API
|
||||
- Async request handling
|
||||
- Multi-stage processing support
|
||||
|
||||
## Installation and Usage
|
||||
|
||||
### Installation
|
||||
```bash
|
||||
# Install vLLM first
|
||||
pip install vllm>=0.2.0
|
||||
|
||||
# Install vLLM-omni
|
||||
pip install vllm-omni
|
||||
|
||||
# Or install from source
|
||||
git clone https://github.com/hsliuustc0106/vllm-omni
|
||||
cd vllm-omni
|
||||
pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
### Basic Usage
|
||||
```python
|
||||
from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
|
||||
|
||||
# Create stage configurations
|
||||
ar_config = create_ar_stage_config(
|
||||
stage_id=0,
|
||||
model_path="Qwen/Qwen3-0.6B",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"]
|
||||
)
|
||||
|
||||
dit_config = create_dit_stage_config(
|
||||
stage_id=1,
|
||||
model_path="stabilityai/stable-diffusion-2-1",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"]
|
||||
)
|
||||
|
||||
# Create OmniLLM instance
|
||||
omni_llm = OmniLLM([ar_config, dit_config])
|
||||
|
||||
# Generate with multi-stage processing
|
||||
stage_args = [
|
||||
{"prompt": "A beautiful landscape"},
|
||||
{"prompt": "A beautiful landscape"}
|
||||
]
|
||||
|
||||
outputs = omni_llm.generate(stage_args)
|
||||
```
|
||||
|
||||
## Next Steps
|
||||
|
||||
### Immediate Priorities
|
||||
1. **Complete Model Runners**: Implement OmniDiffusionModelRunner and OmniARModelRunner
|
||||
2. **Executor Integration**: Complete executor implementations and vLLM integration
|
||||
3. **Testing**: Implement comprehensive test suite
|
||||
4. **Documentation**: Create user guides and API documentation
|
||||
|
||||
### Medium-term Goals
|
||||
1. **Performance Optimization**: Optimize memory usage and inference speed
|
||||
2. **Advanced Features**: Implement advanced scheduling and caching strategies
|
||||
3. **Model Support**: Add support for more model architectures
|
||||
4. **Production Readiness**: Add monitoring, logging, and deployment tools
|
||||
|
||||
### Long-term Vision
|
||||
1. **Multi-GPU Support**: Optimize for distributed inference
|
||||
2. **Custom Architectures**: Support for custom model architectures
|
||||
3. **Advanced Multimodal**: Enhanced multimodal fusion capabilities
|
||||
4. **Ecosystem Integration**: Integration with popular ML frameworks
|
||||
|
||||
## Technical Considerations
|
||||
|
||||
### vLLM Compatibility
|
||||
- Maintains compatibility with vLLM V1 architecture
|
||||
- Reuses proven components (scheduler, executor, worker patterns)
|
||||
- Minimal modifications to existing vLLM codebase
|
||||
|
||||
### Memory Management
|
||||
- Efficient caching system for DiT models
|
||||
- Memory sharing between stages
|
||||
- Garbage collection optimization
|
||||
|
||||
### Extensibility
|
||||
- Plugin-based architecture
|
||||
- Easy integration of new modalities
|
||||
- Configurable stage processing
|
||||
|
||||
### Performance
|
||||
- Optimized for both AR and DiT models
|
||||
- Efficient multi-stage processing
|
||||
- Scalable to multiple GPUs
|
||||
|
||||
## Conclusion
|
||||
|
||||
The vLLM-omni implementation provides a solid foundation for multi-modality model inference with non-autoregressive structures. The current implementation covers the core functionality needed for basic multi-stage processing, with a clear path forward for advanced features and optimizations.
|
||||
|
||||
The modular architecture ensures maintainability and extensibility, while the vLLM integration provides a proven foundation for production deployment. The comprehensive testing strategy and documentation ensure reliability and ease of use.
|
||||
|
||||
This implementation represents a significant step forward in making advanced multimodal AI models accessible and efficient for production use cases.
|
||||
+17
-18
@@ -100,9 +100,9 @@ class TestOmniStageConfig:
|
||||
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
|
||||
num_inference_steps=50,
|
||||
guidance_scale=7.5,
|
||||
use_diffusers=False,
|
||||
)
|
||||
config = OmniStageConfig(
|
||||
stage_id=1,
|
||||
@@ -130,27 +130,24 @@ class TestOmniStageConfig:
|
||||
#### test_dit_config.py
|
||||
```python
|
||||
import pytest
|
||||
from vllm_omni.config.dit_config import DiTConfig
|
||||
from vllm.config import ModelConfig
|
||||
from vllm_omni.config import DiTConfig
|
||||
|
||||
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.model_config is None
|
||||
assert config.use_diffusers is False
|
||||
assert config.diffusers_pipeline is None
|
||||
|
||||
def test_diffusers_config(self):
|
||||
"""Test DiT configuration with diffusers"""
|
||||
config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
model_config=ModelConfig(model="stabilityai/stable-diffusion-2-1"),
|
||||
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"
|
||||
@@ -431,6 +428,7 @@ class TestCLIIntegration:
|
||||
#### test_full_pipeline.py
|
||||
```python
|
||||
import pytest
|
||||
from vllm.config import ModelConfig
|
||||
from vllm_omni.core.omni_llm import OmniLLM
|
||||
from vllm_omni.config.stage_config import OmniStageConfig, DiTConfig
|
||||
|
||||
@@ -454,9 +452,10 @@ class TestFullPipeline:
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"],
|
||||
dit_config=DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=10 # Reduced for testing
|
||||
model_config=ModelConfig(model="stabilityai/stable-diffusion-2-1"),
|
||||
num_inference_steps=10,
|
||||
guidance_scale=7.5,
|
||||
use_diffusers=False,
|
||||
)
|
||||
)
|
||||
]
|
||||
@@ -631,9 +630,9 @@ def sample_stage_configs():
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"],
|
||||
dit_config=DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=10
|
||||
model_config=ModelConfig(model="test-dit-model"),
|
||||
num_inference_steps=10,
|
||||
guidance_scale=7.5,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
+11
-44
@@ -1,49 +1,16 @@
|
||||
# vLLM-omni Examples
|
||||
|
||||
This directory contains examples demonstrating how to use vLLM-omni for multi-modality model inference and serving.
|
||||
A curated set of runnable samples demonstrating vLLM-omni usage.
|
||||
|
||||
## Directory Structure
|
||||
## Directory Overview
|
||||
|
||||
- `basic/` - Simple usage examples and API client
|
||||
- `advanced/` - Advanced features and configurations
|
||||
- `multimodal/` - Multi-modality specific examples
|
||||
- **basic/** – Quick-start scripts for text-only flows.
|
||||
- `text_generation.py` – Minimal CPU-friendly text generation with vLLM.
|
||||
- `api_client.py` – Calls the FastAPI server using the JSON API.
|
||||
- `docker_run.py` – Template for launching the server in Docker.
|
||||
- **omni/** – Multi-stage AR → DiT pipelines powered by diffusers.
|
||||
- `ar_dit_diffusers.py` – Runs Qwen3 (AR) feeding Stable Diffusion 2.1 (DiT).
|
||||
- `configs/ar_dit_local.yaml` – Sample plugin config referencing local weights.
|
||||
- **advanced/** – Placeholder for upcoming complex flows (Ray, benchmarking).
|
||||
|
||||
## Quick Start
|
||||
|
||||
1. **Basic Usage:**
|
||||
```bash
|
||||
cd examples/basic
|
||||
python simple_usage.py
|
||||
```
|
||||
|
||||
2. **API Server:**
|
||||
```bash
|
||||
# Start server
|
||||
vllm serve Qwen/Qwen3-0.6B --omni --port 8000
|
||||
|
||||
# Test with client
|
||||
python api_client.py
|
||||
```
|
||||
|
||||
## Examples Overview
|
||||
|
||||
### Basic Examples (`basic/`)
|
||||
- `simple_usage.py` - Direct library usage (sync/async)
|
||||
- `api_client.py` - HTTP API client examples
|
||||
- `README.md` - Basic examples documentation
|
||||
|
||||
### Advanced Examples (`advanced/`)
|
||||
- Multi-stage processing
|
||||
- Custom configurations
|
||||
- Performance optimization
|
||||
|
||||
### Multi-modality Examples (`multimodal/`)
|
||||
- Text-to-image generation
|
||||
- Image-to-text processing
|
||||
- Multi-modal pipelines
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.8+
|
||||
- vLLM-omni installed
|
||||
- For API examples: HTTP server running
|
||||
See the subdirectory READMEs for usage instructions.
|
||||
@@ -1,76 +1,10 @@
|
||||
# Basic vLLM-omni Examples
|
||||
# Basic Examples
|
||||
|
||||
This directory contains simple examples demonstrating how to use vLLM-omni for text generation.
|
||||
Quick-start scripts for experimenting with vLLM-omni.
|
||||
|
||||
## Examples
|
||||
## Scripts
|
||||
|
||||
### 1. `simple_usage.py` - Direct Library Usage
|
||||
|
||||
Demonstrates how to use vLLM-omni directly in your Python code:
|
||||
|
||||
```bash
|
||||
# Run the example
|
||||
python simple_usage.py
|
||||
```
|
||||
|
||||
**Features:**
|
||||
- Synchronous text generation
|
||||
- Asynchronous text generation
|
||||
- Multi-prompt processing
|
||||
- Shows both text output and token information
|
||||
|
||||
### 2. `api_client.py` - API Server Usage
|
||||
|
||||
Demonstrates how to interact with the vLLM-omni API server:
|
||||
|
||||
```bash
|
||||
# First, start the server
|
||||
vllm serve Qwen/Qwen3-0.6B --omni --port 8000
|
||||
|
||||
# Then run the client (in another terminal)
|
||||
python api_client.py
|
||||
```
|
||||
|
||||
**Features:**
|
||||
- Health check
|
||||
- Text generation via HTTP API
|
||||
- Multiple prompt types
|
||||
- Full response structure inspection
|
||||
|
||||
## Quick Start
|
||||
|
||||
1. **Install vLLM-omni:**
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
2. **Run direct usage example:**
|
||||
```bash
|
||||
cd examples/basic
|
||||
python simple_usage.py
|
||||
```
|
||||
|
||||
3. **Run API server example:**
|
||||
```bash
|
||||
# Terminal 1: Start server
|
||||
vllm serve Qwen/Qwen3-0.6B --omni --port 8000
|
||||
|
||||
# Terminal 2: Run client
|
||||
cd examples/basic
|
||||
python api_client.py
|
||||
```
|
||||
|
||||
## Expected Output
|
||||
|
||||
The examples will show:
|
||||
- ✅ Generated text responses
|
||||
- ✅ Token counts and completion status
|
||||
- ✅ Stage processing information
|
||||
- ✅ Error handling
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.8+
|
||||
- vLLM-omni installed
|
||||
- For API examples: HTTP server running
|
||||
- `text_generation.py` – Minimal standalone usage showing how to create an LLM and generate text locally.
|
||||
- `api_client.py` – Demonstrates calling the FastAPI `/generate` endpoint from Python.
|
||||
|
||||
Run any script with `python examples/basic/<script>.py`. Ensure the `vllm_omni` conda environment is active and dependencies are installed.
|
||||
@@ -1,178 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Docker example for vLLM-omni.
|
||||
|
||||
This example shows how to test the vLLM-omni API server running in Docker.
|
||||
"""
|
||||
|
||||
import requests
|
||||
import json
|
||||
import time
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def check_docker_running():
|
||||
"""Check if Docker is running."""
|
||||
try:
|
||||
result = subprocess.run(['docker', 'info'],
|
||||
capture_output=True, text=True, timeout=5)
|
||||
return result.returncode == 0
|
||||
except (subprocess.TimeoutExpired, FileNotFoundError):
|
||||
return False
|
||||
|
||||
|
||||
def start_docker_server():
|
||||
"""Start the vLLM-omni server in Docker."""
|
||||
print("🐳 Starting vLLM-omni server in Docker...")
|
||||
|
||||
# Check if container is already running
|
||||
result = subprocess.run(['docker', 'ps', '--filter', 'name=vllm-omni-cpu', '--format', '{{.Names}}'],
|
||||
capture_output=True, text=True)
|
||||
|
||||
if 'vllm-omni-cpu' in result.stdout:
|
||||
print("✅ Server already running")
|
||||
return True
|
||||
|
||||
# Start the server
|
||||
cmd = [
|
||||
'docker', 'run', '-d', '--name', 'vllm-omni-cpu',
|
||||
'-p', '8000:8000',
|
||||
'vllm-omni-cpu:latest',
|
||||
'vllm', 'serve', 'Qwen/Qwen3-0.6B', '--omni', '--port', '8000'
|
||||
]
|
||||
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
print("✅ Server started successfully")
|
||||
|
||||
# Wait for server to be ready
|
||||
print("⏳ Waiting for server to be ready...")
|
||||
for i in range(30): # Wait up to 30 seconds
|
||||
if test_server_health():
|
||||
print("✅ Server is ready!")
|
||||
return True
|
||||
time.sleep(1)
|
||||
|
||||
print("❌ Server failed to start within 30 seconds")
|
||||
return False
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"❌ Failed to start server: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_server_health():
|
||||
"""Test if the server is healthy."""
|
||||
try:
|
||||
response = requests.get("http://localhost:8000/health", timeout=2)
|
||||
return response.status_code == 200
|
||||
except requests.exceptions.RequestException:
|
||||
return False
|
||||
|
||||
|
||||
def test_generation():
|
||||
"""Test text generation via API."""
|
||||
print("\n🧪 Testing text generation...")
|
||||
|
||||
test_cases = [
|
||||
{
|
||||
"prompt": "Hello, how are you?",
|
||||
"max_tokens": 50,
|
||||
"temperature": 0.7
|
||||
},
|
||||
{
|
||||
"prompt": "Tell me a joke",
|
||||
"max_tokens": 100,
|
||||
"temperature": 0.9
|
||||
},
|
||||
{
|
||||
"prompt": "What is artificial intelligence?",
|
||||
"max_tokens": 150,
|
||||
"temperature": 0.5
|
||||
}
|
||||
]
|
||||
|
||||
for i, test_case in enumerate(test_cases, 1):
|
||||
print(f"\nTest {i}: {test_case['prompt']}")
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
"http://localhost:8000/generate",
|
||||
json={
|
||||
"prompts": [test_case["prompt"]],
|
||||
"max_tokens": test_case["max_tokens"],
|
||||
"temperature": test_case["temperature"]
|
||||
},
|
||||
timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
text = result["outputs"][0]["text"]
|
||||
print(f"✅ Response: {text}")
|
||||
else:
|
||||
print(f"❌ Error: {response.status_code} - {response.text}")
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"❌ Request failed: {e}")
|
||||
|
||||
|
||||
def cleanup():
|
||||
"""Clean up Docker container."""
|
||||
print("\n🧹 Cleaning up...")
|
||||
try:
|
||||
subprocess.run(['docker', 'stop', 'vllm-omni-cpu'],
|
||||
capture_output=True, text=True)
|
||||
subprocess.run(['docker', 'rm', 'vllm-omni-cpu'],
|
||||
capture_output=True, text=True)
|
||||
print("✅ Cleanup completed")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Cleanup warning: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main example function."""
|
||||
print("🐳 vLLM-omni Docker Example")
|
||||
print("=" * 40)
|
||||
|
||||
# Check Docker
|
||||
if not check_docker_running():
|
||||
print("❌ Docker is not running. Please start Docker Desktop.")
|
||||
sys.exit(1)
|
||||
|
||||
print("✅ Docker is running")
|
||||
|
||||
# Check if image exists
|
||||
result = subprocess.run(['docker', 'images', 'vllm-omni-cpu:latest'],
|
||||
capture_output=True, text=True)
|
||||
if 'vllm-omni-cpu' not in result.stdout:
|
||||
print("❌ Docker image not found. Please build it first:")
|
||||
print(" docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu .")
|
||||
sys.exit(1)
|
||||
|
||||
print("✅ Docker image found")
|
||||
|
||||
try:
|
||||
# Start server
|
||||
if not start_docker_server():
|
||||
sys.exit(1)
|
||||
|
||||
# Test generation
|
||||
test_generation()
|
||||
|
||||
print("\n🎉 Docker example completed successfully!")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n⏹️ Example interrupted by user")
|
||||
finally:
|
||||
# Ask if user wants to keep server running
|
||||
keep_running = input("\n🤔 Keep the server running? (y/N): ").lower().strip()
|
||||
if keep_running != 'y':
|
||||
cleanup()
|
||||
else:
|
||||
print("✅ Server kept running. Access it at http://localhost:8000")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File renamed without changes.
@@ -0,0 +1,31 @@
|
||||
# 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.).
|
||||
@@ -0,0 +1,175 @@
|
||||
"""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 (
|
||||
OmniLLM,
|
||||
DiTConfig,
|
||||
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()
|
||||
@@ -0,0 +1,28 @@
|
||||
# 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
|
||||
+1
-3
@@ -34,6 +34,7 @@ dependencies = [
|
||||
"numpy>=1.21.0",
|
||||
"pillow>=8.0.0",
|
||||
"opencv-python>=4.5.0",
|
||||
"PyYAML>=6.0",
|
||||
"gradio>=3.0.0",
|
||||
"fastapi>=0.100.0",
|
||||
"uvicorn>=0.20.0",
|
||||
@@ -74,9 +75,6 @@ Documentation = "https://vllm-omni.readthedocs.io"
|
||||
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"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["vllm_omni*"]
|
||||
|
||||
@@ -5,6 +5,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
|
||||
|
||||
@@ -1,83 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# vLLM-omni Docker Setup Script for macOS
|
||||
# This script helps set up Docker for vLLM-omni on macOS
|
||||
|
||||
set -e
|
||||
|
||||
echo "🐳 vLLM-omni Docker Setup for macOS"
|
||||
echo "=================================="
|
||||
|
||||
# Check if Docker is installed
|
||||
if ! command -v docker &> /dev/null; then
|
||||
echo "❌ Docker is not installed. Please install Docker Desktop for Mac:"
|
||||
echo " https://www.docker.com/products/docker-desktop/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if Docker is running
|
||||
if ! docker info &> /dev/null; then
|
||||
echo "❌ Docker is not running. Please start Docker Desktop."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "✅ Docker is installed and running"
|
||||
|
||||
# Check if NVIDIA Docker runtime is available (optional for macOS)
|
||||
if docker info | grep -q nvidia; then
|
||||
echo "✅ NVIDIA Docker runtime detected"
|
||||
GPU_SUPPORT=true
|
||||
else
|
||||
echo "⚠️ NVIDIA Docker runtime not detected (normal on macOS)"
|
||||
GPU_SUPPORT=false
|
||||
fi
|
||||
|
||||
# Create necessary directories
|
||||
echo "📁 Creating directories..."
|
||||
mkdir -p models logs notebooks
|
||||
|
||||
# Build the Docker image
|
||||
echo "🔨 Building vLLM-omni Docker image..."
|
||||
docker build -f docker/Dockerfile.cpu -t vllm-omni-cpu:latest .
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "✅ Docker image built successfully"
|
||||
else
|
||||
echo "❌ Failed to build Docker image"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Test the image
|
||||
echo "🧪 Testing Docker image..."
|
||||
docker run --rm vllm-omni-cpu:latest python -c "import vllm_omni; print('vLLM-omni imported successfully')"
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "✅ Docker image test passed"
|
||||
else
|
||||
echo "❌ Docker image test failed"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "🎉 Setup complete! You can now use:"
|
||||
echo ""
|
||||
echo " # Start the server:"
|
||||
echo " docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu"
|
||||
echo ""
|
||||
echo " # Start with Jupyter:"
|
||||
echo " docker-compose -f docker/docker-compose.cpu.yml --profile jupyter up"
|
||||
echo ""
|
||||
echo " # Run a single command:"
|
||||
echo " docker run --rm -p 8000:8000 vllm-omni-cpu:latest vllm serve Qwen/Qwen3-0.6B --omni --port 8000"
|
||||
echo ""
|
||||
echo " # Interactive shell:"
|
||||
echo " docker run -it --rm vllm-omni-cpu:latest bash"
|
||||
echo ""
|
||||
|
||||
# Optional: Start the server
|
||||
read -p "🚀 Start the vLLM-omni server now? (y/N): " -n 1 -r
|
||||
echo
|
||||
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
||||
echo "Starting vLLM-omni server..."
|
||||
docker-compose -f docker/docker-compose.cpu.yml up vllm-omni-cpu
|
||||
fi
|
||||
@@ -1,147 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Model download script for vLLM-omni
|
||||
Downloads commonly used AR and DiT models
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
def download_ar_models(models=None):
|
||||
"""Download AR (Autoregressive) models"""
|
||||
if models is None:
|
||||
models = [
|
||||
"Qwen/Qwen3-0.6B",
|
||||
]
|
||||
|
||||
print("🚀 Downloading AR models...")
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
except ImportError:
|
||||
print("❌ transformers not installed. Installing...")
|
||||
os.system("pip install transformers")
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
for model_name in models:
|
||||
print(f"\n📥 Downloading AR model: {model_name}")
|
||||
try:
|
||||
print(f" - Downloading tokenizer...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
print(f" - Downloading model...")
|
||||
model = AutoModel.from_pretrained(model_name)
|
||||
|
||||
print(f"✅ {model_name} downloaded successfully")
|
||||
|
||||
# Print model info
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
size_mb = total_params * 4 / (1024 * 1024)
|
||||
print(f" 📊 Parameters: {total_params:,} ({size_mb:.1f} MB)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to download {model_name}: {e}")
|
||||
|
||||
def download_dit_models(models=None):
|
||||
"""Download DiT (Diffusion Transformer) models"""
|
||||
if models is None:
|
||||
models = [
|
||||
"stabilityai/stable-diffusion-2-1",
|
||||
"runwayml/stable-diffusion-v1-5"
|
||||
]
|
||||
|
||||
print("\n🚀 Downloading DiT models...")
|
||||
|
||||
try:
|
||||
from diffusers import StableDiffusionPipeline
|
||||
except ImportError:
|
||||
print("❌ diffusers not installed. Installing...")
|
||||
os.system("pip install diffusers")
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
for model_name in models:
|
||||
print(f"\n📥 Downloading DiT model: {model_name}")
|
||||
try:
|
||||
print(f" - Downloading pipeline...")
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_name)
|
||||
|
||||
print(f"✅ {model_name} downloaded successfully")
|
||||
|
||||
# Print model info
|
||||
total_params = sum(p.numel() for p in pipe.unet.parameters())
|
||||
size_mb = total_params * 4 / (1024 * 1024)
|
||||
print(f" 📊 UNet parameters: {total_params:,} ({size_mb:.1f} MB)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to download {model_name}: {e}")
|
||||
|
||||
def check_cache_directory():
|
||||
"""Check and display cache directory info"""
|
||||
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
|
||||
print(f"📁 Cache directory: {cache_dir}")
|
||||
|
||||
if cache_dir.exists():
|
||||
total_size = sum(f.stat().st_size for f in cache_dir.rglob('*') if f.is_file())
|
||||
size_gb = total_size / (1024**3)
|
||||
print(f"📊 Total cache size: {size_gb:.2f} GB")
|
||||
|
||||
# List downloaded models
|
||||
model_dirs = [d for d in cache_dir.iterdir() if d.is_dir() and d.name.startswith("models--")]
|
||||
if model_dirs:
|
||||
print(f"📋 Downloaded models ({len(model_dirs)}):")
|
||||
for model_dir in sorted(model_dirs):
|
||||
model_name = model_dir.name.replace("models--", "").replace("--", "/")
|
||||
print(f" - {model_name}")
|
||||
else:
|
||||
print("📁 Cache directory does not exist yet")
|
||||
|
||||
def clear_cache():
|
||||
"""Clear model cache"""
|
||||
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
|
||||
|
||||
if cache_dir.exists():
|
||||
print(f"🗑️ Clearing cache directory: {cache_dir}")
|
||||
import shutil
|
||||
shutil.rmtree(cache_dir)
|
||||
print("✅ Cache cleared successfully")
|
||||
else:
|
||||
print("📁 Cache directory does not exist")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Download models for vLLM-omni")
|
||||
parser.add_argument("--ar-models", nargs="+", help="AR models to download")
|
||||
parser.add_argument("--dit-models", nargs="+", help="DiT models to download")
|
||||
parser.add_argument("--all", action="store_true", help="Download all default models")
|
||||
parser.add_argument("--check-cache", action="store_true", help="Check cache directory")
|
||||
parser.add_argument("--clear-cache", action="store_true", help="Clear cache directory")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print("🎯 vLLM-omni Model Downloader")
|
||||
print("=" * 40)
|
||||
|
||||
if args.clear_cache:
|
||||
clear_cache()
|
||||
return
|
||||
|
||||
if args.check_cache:
|
||||
check_cache_directory()
|
||||
return
|
||||
|
||||
if args.all or args.ar_models:
|
||||
download_ar_models(args.ar_models)
|
||||
|
||||
if args.all or args.dit_models:
|
||||
download_dit_models(args.dit_models)
|
||||
|
||||
if not any([args.all, args.ar_models, args.dit_models, args.check_cache, args.clear_cache]):
|
||||
print("❓ No action specified. Use --help for options.")
|
||||
print("\nQuick start:")
|
||||
print(" python scripts/download_models.py --all")
|
||||
print(" python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B")
|
||||
print(" python scripts/download_models.py --check-cache")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,82 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Quick vLLM-omni Test Script
|
||||
# Fast test for basic functionality after changes
|
||||
# Usage: ./scripts/quick_test.sh [port]
|
||||
|
||||
set -e
|
||||
|
||||
PORT=${1:-8000}
|
||||
HOST="localhost"
|
||||
|
||||
# Colors
|
||||
GREEN='\033[0;32m'
|
||||
RED='\033[0;31m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m'
|
||||
|
||||
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
|
||||
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
|
||||
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
|
||||
|
||||
# Cleanup
|
||||
cleanup() {
|
||||
log_info "Cleaning up..."
|
||||
pkill -f "vllm serve" 2>/dev/null || true
|
||||
}
|
||||
|
||||
trap cleanup EXIT
|
||||
|
||||
# Quick test
|
||||
main() {
|
||||
log_info "Quick vLLM-omni test on port $PORT"
|
||||
|
||||
# Test imports
|
||||
log_info "Testing imports..."
|
||||
python -c "import vllm_omni; from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM; print('✅ Imports OK')" || {
|
||||
log_error "Import failed"
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Start server
|
||||
log_info "Starting server..."
|
||||
vllm serve ./models/Qwen3-0.6B --omni --port $PORT > /dev/null 2>&1 &
|
||||
SERVER_PID=$!
|
||||
|
||||
# Wait for server
|
||||
sleep 15
|
||||
|
||||
# Test health with retry
|
||||
log_info "Testing health endpoint..."
|
||||
for i in {1..5}; do
|
||||
if curl -s http://$HOST:$PORT/health | grep -q "healthy"; then
|
||||
log_success "Health check passed"
|
||||
break
|
||||
else
|
||||
if [ $i -eq 5 ]; then
|
||||
log_error "Health check failed after 5 attempts"
|
||||
exit 1
|
||||
fi
|
||||
log_info "Health check attempt $i failed, retrying in 3 seconds..."
|
||||
sleep 3
|
||||
fi
|
||||
done
|
||||
|
||||
# Test generation
|
||||
log_info "Testing text generation..."
|
||||
response=$(curl -s -X POST http://$HOST:$PORT/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"prompts": ["Quick test"], "max_tokens": 10, "temperature": 0.7}')
|
||||
|
||||
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 "Generation test passed"
|
||||
else
|
||||
log_error "Generation test failed"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
log_success "Quick test completed successfully!"
|
||||
}
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
main "$@"
|
||||
+74
-12
@@ -122,16 +122,22 @@ start_server() {
|
||||
# Test health endpoint
|
||||
test_health() {
|
||||
log_info "Testing health endpoint..."
|
||||
local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
|
||||
local http_code="${response: -3}"
|
||||
local body="${response%???}"
|
||||
|
||||
if [ "$http_code" = "200" ]; then
|
||||
log_success "Health check passed: $body"
|
||||
else
|
||||
log_error "Health check failed (HTTP $http_code): $body"
|
||||
return 1
|
||||
fi
|
||||
for attempt in {1..10}; do
|
||||
local response=$(curl -s -w "%{http_code}" http://$HOST:$PORT/health 2>/dev/null || echo "000")
|
||||
local http_code="${response: -3}"
|
||||
local body="${response%???}"
|
||||
|
||||
if [ "$http_code" = "200" ]; then
|
||||
log_success "Health check passed: $body"
|
||||
return 0
|
||||
fi
|
||||
|
||||
log_warning "Health check attempt $attempt failed (HTTP $http_code). Retrying in 3s..."
|
||||
sleep 3
|
||||
done
|
||||
|
||||
log_error "Health check failed after multiple attempts"
|
||||
return 1
|
||||
}
|
||||
|
||||
# Test info endpoint
|
||||
@@ -186,6 +192,60 @@ test_api_client() {
|
||||
fi
|
||||
}
|
||||
|
||||
# Test AR → DiT pipeline example
|
||||
test_ar_dit_pipeline() {
|
||||
log_info "Testing AR → DiT diffusers pipeline example..."
|
||||
|
||||
if ! python - 2>/dev/null <<'PY'
|
||||
import importlib
|
||||
for pkg in ("torch", "diffusers"):
|
||||
if importlib.util.find_spec(pkg) is None:
|
||||
raise ImportError(pkg)
|
||||
PY
|
||||
then
|
||||
log_warning "Required packages for diffusers pipeline not found, skipping AR → DiT test"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [ ! -f "examples/omni/ar_dit_diffusers.py" ]; then
|
||||
log_warning "AR → DiT example not found, skipping"
|
||||
return 0
|
||||
fi
|
||||
|
||||
local output_path="logs/ar_dit_pipeline.png"
|
||||
local log_path="logs/ar_dit_pipeline.log"
|
||||
rm -f "$output_path"
|
||||
|
||||
local cmd=(python examples/omni/ar_dit_diffusers.py
|
||||
--prompt "Smoke test prompt of a lighthouse"
|
||||
--seed 0
|
||||
--output "$output_path")
|
||||
|
||||
local env_prefix=()
|
||||
if [ -n "$AR_MODEL" ]; then
|
||||
env_prefix+=("AR_MODEL=$AR_MODEL")
|
||||
fi
|
||||
if [ -n "$DIT_MODEL" ]; then
|
||||
env_prefix+=("DIT_MODEL=$DIT_MODEL")
|
||||
fi
|
||||
|
||||
if ! env "${env_prefix[@]}" "${cmd[@]}" >"$log_path" 2>&1; then
|
||||
log_error "AR → DiT pipeline test failed"
|
||||
log_info "AR → DiT logs:"
|
||||
cat "$log_path"
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [ ! -f "$output_path" ]; then
|
||||
log_error "AR → DiT pipeline did not produce an image"
|
||||
log_info "AR → DiT logs:"
|
||||
cat "$log_path"
|
||||
return 1
|
||||
fi
|
||||
|
||||
log_success "AR → DiT pipeline test passed (output: $output_path)"
|
||||
}
|
||||
|
||||
# Test simple usage example
|
||||
test_simple_usage() {
|
||||
log_info "Testing simple usage example..."
|
||||
@@ -250,10 +310,11 @@ run_tests() {
|
||||
test_generation
|
||||
test_performance
|
||||
test_api_client
|
||||
|
||||
test_ar_dit_pipeline
|
||||
|
||||
# Note: Simple usage test is commented out as it starts its own server
|
||||
# test_simple_usage
|
||||
|
||||
|
||||
log_success "All tests completed successfully!"
|
||||
echo "=========================================="
|
||||
log_info "Test Summary:"
|
||||
@@ -262,6 +323,7 @@ run_tests() {
|
||||
log_info "✅ Text generation functionality"
|
||||
log_info "✅ Performance metrics"
|
||||
log_info "✅ API client integration"
|
||||
log_info "✅ AR → DiT diffusers pipeline"
|
||||
log_info "✅ All imports working correctly"
|
||||
}
|
||||
|
||||
|
||||
@@ -3,25 +3,16 @@ Stage configuration for vLLM-omni multi-stage processing.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Type, Literal, Any
|
||||
from vllm.config import VllmConfig
|
||||
from typing import List, Optional, Type, Literal, Any, Dict
|
||||
from vllm.config import (
|
||||
VllmConfig,
|
||||
DeviceConfig,
|
||||
LoadConfig,
|
||||
ModelConfig,
|
||||
CompilationConfig,
|
||||
)
|
||||
from vllm.executor.executor_base import ExecutorBase as Executor
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiTConfig:
|
||||
"""Configuration for DiT (Diffusion Transformer) stages."""
|
||||
model_type: str
|
||||
scheduler_type: str
|
||||
num_inference_steps: int
|
||||
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 DiTCacheTensor:
|
||||
"""Configuration for DiT cache tensors."""
|
||||
@@ -41,18 +32,46 @@ class DiTCacheConfig:
|
||||
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
|
||||
executor_class: Type[Executor] = None # Will be set based on engine_type
|
||||
dit_config: Optional[DiTConfig] = None
|
||||
cache_config: Optional[DiTCacheConfig] = 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):
|
||||
@@ -94,7 +113,8 @@ def create_ar_stage_config(
|
||||
input_modalities: List[str] = None,
|
||||
output_modalities: List[str] = None,
|
||||
vllm_config: Optional[VllmConfig] = None,
|
||||
executor_class: Optional[Type[Executor]] = 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:
|
||||
@@ -111,19 +131,19 @@ def create_ar_stage_config(
|
||||
input_modalities=input_modalities,
|
||||
output_modalities=output_modalities,
|
||||
vllm_config=vllm_config,
|
||||
executor_class=executor_class
|
||||
executor_class=executor_class,
|
||||
default_stage_args=default_stage_args,
|
||||
)
|
||||
|
||||
|
||||
def create_dit_stage_config(
|
||||
stage_id: int,
|
||||
model_path: str,
|
||||
input_modalities: List[str] = None,
|
||||
output_modalities: List[str] = None,
|
||||
input_modalities: Optional[List[str]] = None,
|
||||
output_modalities: Optional[List[str]] = None,
|
||||
dit_config: Optional[DiTConfig] = None,
|
||||
cache_config: Optional[DiTCacheConfig] = None,
|
||||
vllm_config: Optional[VllmConfig] = None,
|
||||
executor_class: Optional[Type[Executor]] = 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:
|
||||
@@ -132,22 +152,15 @@ def create_dit_stage_config(
|
||||
output_modalities = ["image"]
|
||||
|
||||
if dit_config is None:
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=50
|
||||
)
|
||||
|
||||
# For now, we'll create minimal configs
|
||||
# vllm_config and executor_class will be handled by the LLM class
|
||||
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,
|
||||
vllm_config=vllm_config,
|
||||
dit_config=dit_config,
|
||||
cache_config=cache_config,
|
||||
executor_class=executor_class
|
||||
executor_class=executor_class,
|
||||
default_stage_args=default_stage_args,
|
||||
)
|
||||
@@ -1 +0,0 @@
|
||||
from vllm.v1.core.sched import SchedulerInterface
|
||||
@@ -1,229 +0,0 @@
|
||||
"""
|
||||
Diffusion scheduler for DiT (Diffusion Transformer) models in vLLM-omni.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Any, TYPE_CHECKING
|
||||
from vllm.v1.core.sched.scheduler import SchedulerInterface
|
||||
from vllm.config import VllmConfig, KVCacheConfig
|
||||
from vllm.v1.core.sched.structured_output_manager import StructuredOutputManager
|
||||
from vllm.v1.core.sched.mm_registry import MultiModalRegistry, MULTIMODAL_REGISTRY
|
||||
|
||||
from ..dit_cache_manager import DiTCacheManager
|
||||
from ...config import DiTCacheConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
||||
|
||||
|
||||
class OmniDiffusionScheduler(SchedulerInterface):
|
||||
"""Scheduler for DiT models with caching optimization."""
|
||||
|
||||
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:
|
||||
super().__init__(
|
||||
vllm_config=vllm_config,
|
||||
kv_cache_config=kv_cache_config,
|
||||
structured_output_manager=structured_output_manager,
|
||||
mm_registry=mm_registry,
|
||||
include_finished_set=include_finished_set,
|
||||
log_stats=log_stats,
|
||||
)
|
||||
|
||||
self.dit_cache_config = dit_cache_config
|
||||
self.dit_cache_manager = DiTCacheManager(dit_cache_config)
|
||||
|
||||
# DiT-specific scheduling state
|
||||
self.dit_requests: Dict[str, Dict[str, Any]] = {}
|
||||
self.current_step: int = 0
|
||||
self.max_steps: int = 50 # Default, will be updated from config
|
||||
|
||||
def schedule(self) -> "SchedulerOutput":
|
||||
"""Schedule DiT requests with caching optimization."""
|
||||
# Get pending requests
|
||||
pending_requests = self._get_pending_requests()
|
||||
|
||||
if not pending_requests:
|
||||
return self._create_empty_scheduler_output()
|
||||
|
||||
# Apply DiT-specific scheduling logic
|
||||
scheduled_requests = self._schedule_dit_requests(pending_requests)
|
||||
|
||||
# Create scheduler output
|
||||
return self._create_scheduler_output(scheduled_requests)
|
||||
|
||||
def _get_pending_requests(self) -> List[Dict[str, Any]]:
|
||||
"""Get pending DiT requests."""
|
||||
# This would integrate with vLLM's request management system
|
||||
# For now, we'll return a mock implementation
|
||||
return []
|
||||
|
||||
def _schedule_dit_requests(self, requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Apply DiT-specific scheduling logic."""
|
||||
scheduled_requests = []
|
||||
|
||||
for request in requests:
|
||||
# Check cache for this request
|
||||
cached_result = self.dit_cache_manager.get_cache(request['request_id'])
|
||||
|
||||
if cached_result is not None:
|
||||
# Use cached result
|
||||
request['cached'] = True
|
||||
request['cached_result'] = cached_result
|
||||
else:
|
||||
# Allocate cache for new request
|
||||
cache_tensor = self.dit_cache_manager.allocate_cache(
|
||||
request['request_id'],
|
||||
request.get('cache_size', 1024)
|
||||
)
|
||||
request['cache_tensor'] = cache_tensor
|
||||
request['cached'] = False
|
||||
|
||||
# Apply DiT-specific scheduling policies
|
||||
request = self._apply_dit_scheduling_policies(request)
|
||||
scheduled_requests.append(request)
|
||||
|
||||
return scheduled_requests
|
||||
|
||||
def _apply_dit_scheduling_policies(self, request: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Apply DiT-specific scheduling policies."""
|
||||
# For now, we'll use a simple FIFO policy
|
||||
# In practice, this could include:
|
||||
# - Priority-based scheduling
|
||||
# - Batch size optimization
|
||||
# - Memory-aware scheduling
|
||||
# - Quality vs speed trade-offs
|
||||
|
||||
request['priority'] = request.get('priority', 0)
|
||||
request['batch_size'] = request.get('batch_size', 1)
|
||||
request['scheduling_policy'] = 'fifo'
|
||||
|
||||
return request
|
||||
|
||||
def _create_empty_scheduler_output(self) -> "SchedulerOutput":
|
||||
"""Create an empty scheduler output."""
|
||||
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
||||
|
||||
return SchedulerOutput(
|
||||
scheduled_seq_groups=[],
|
||||
ignored_seq_groups=[],
|
||||
preempted_seq_groups=[],
|
||||
num_preemption_groups=0,
|
||||
num_batched_tokens=0,
|
||||
blocks_to_swap_in={},
|
||||
blocks_to_swap_out={},
|
||||
blocks_to_copy={},
|
||||
ignored_seq_groups=[],
|
||||
num_ignored_seq_groups=0,
|
||||
)
|
||||
|
||||
def _create_scheduler_output(self, scheduled_requests: List[Dict[str, Any]]) -> "SchedulerOutput":
|
||||
"""Create scheduler output from scheduled requests."""
|
||||
from vllm.v1.core.sched.scheduler import SchedulerOutput
|
||||
|
||||
# Convert scheduled requests to vLLM's expected format
|
||||
scheduled_seq_groups = []
|
||||
for request in scheduled_requests:
|
||||
seq_group = self._create_seq_group_from_request(request)
|
||||
scheduled_seq_groups.append(seq_group)
|
||||
|
||||
return SchedulerOutput(
|
||||
scheduled_seq_groups=scheduled_seq_groups,
|
||||
ignored_seq_groups=[],
|
||||
preempted_seq_groups=[],
|
||||
num_preemption_groups=0,
|
||||
num_batched_tokens=sum(req.get('num_tokens', 0) for req in scheduled_requests),
|
||||
blocks_to_swap_in={},
|
||||
blocks_to_swap_out={},
|
||||
blocks_to_copy={},
|
||||
ignored_seq_groups=[],
|
||||
num_ignored_seq_groups=0,
|
||||
)
|
||||
|
||||
def _create_seq_group_from_request(self, request: Dict[str, Any]) -> Any:
|
||||
"""Create a sequence group from a DiT request."""
|
||||
# This would create a proper sequence group
|
||||
# For now, we'll return a mock implementation
|
||||
from vllm.v1.core.sched.sequence import SequenceGroup
|
||||
|
||||
# Mock sequence group creation
|
||||
# In practice, this would properly create a SequenceGroup
|
||||
# with the appropriate metadata for DiT processing
|
||||
return None
|
||||
|
||||
def update_from_output(
|
||||
self,
|
||||
scheduler_output: "SchedulerOutput",
|
||||
model_output: Any,
|
||||
) -> List[Any]:
|
||||
"""Update scheduler state from model output."""
|
||||
# Update cache with new results
|
||||
for seq_group in scheduler_output.scheduled_seq_groups:
|
||||
if hasattr(seq_group, 'request_id'):
|
||||
# Store result in cache
|
||||
self.dit_cache_manager.store_cache(
|
||||
seq_group.request_id,
|
||||
model_output
|
||||
)
|
||||
|
||||
# Update DiT-specific state
|
||||
self.current_step += 1
|
||||
|
||||
# Call parent update method
|
||||
return super().update_from_output(scheduler_output, model_output)
|
||||
|
||||
def add_request(self, request_id: str, **kwargs) -> None:
|
||||
"""Add a new DiT request to the scheduler."""
|
||||
# Store request metadata
|
||||
self.dit_requests[request_id] = {
|
||||
'request_id': request_id,
|
||||
'added_time': self.current_step,
|
||||
'status': 'pending',
|
||||
**kwargs
|
||||
}
|
||||
|
||||
# Call parent add_request method
|
||||
super().add_request(request_id, **kwargs)
|
||||
|
||||
def remove_request(self, request_id: str) -> None:
|
||||
"""Remove a DiT request from the scheduler."""
|
||||
# Clean up request metadata
|
||||
if request_id in self.dit_requests:
|
||||
del self.dit_requests[request_id]
|
||||
|
||||
# Release cache for this request
|
||||
self.dit_cache_manager.release_cache(request_id)
|
||||
|
||||
# Call parent remove_request method
|
||||
super().remove_request(request_id)
|
||||
|
||||
def get_dit_request_info(self, request_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get information about a specific DiT request."""
|
||||
return self.dit_requests.get(request_id)
|
||||
|
||||
def get_cache_statistics(self) -> Dict[str, Any]:
|
||||
"""Get cache statistics for monitoring."""
|
||||
return self.dit_cache_manager.get_statistics()
|
||||
|
||||
def clear_expired_cache(self) -> None:
|
||||
"""Clear expired cache entries."""
|
||||
self.dit_cache_manager.clear_expired_cache()
|
||||
|
||||
def set_max_steps(self, max_steps: int) -> None:
|
||||
"""Set the maximum number of diffusion steps."""
|
||||
self.max_steps = max_steps
|
||||
|
||||
def get_current_step(self) -> int:
|
||||
"""Get the current diffusion step."""
|
||||
return self.current_step
|
||||
|
||||
def reset_step_counter(self) -> None:
|
||||
"""Reset the diffusion step counter."""
|
||||
self.current_step = 0
|
||||
@@ -2,9 +2,10 @@
|
||||
Stage manager for orchestrating multiple engines in vLLM-omni.
|
||||
"""
|
||||
|
||||
from typing import List, Optional
|
||||
from typing import List, Optional, Union
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from ..engine.diffusion_engine import DiffusionEngine
|
||||
|
||||
from ..config import OmniStageConfig
|
||||
|
||||
@@ -15,7 +16,7 @@ class StageManager:
|
||||
def __init__(self, stage_configs: List[OmniStageConfig], log_stats: bool = False):
|
||||
self.stage_configs = stage_configs
|
||||
self.log_stats = log_stats
|
||||
self.engine_list: List[LLMEngine] = []
|
||||
self.engine_list: List[Union[LLMEngine, DiffusionEngine]] = []
|
||||
self.async_engine_list: List[AsyncLLM] = []
|
||||
self._initialized = False
|
||||
self._async_initialized = False
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..config import DiTConfig
|
||||
from ..worker.gpu_diffusion_worker import (
|
||||
DiffusionGPUWorker,
|
||||
DiffusionRunnerOutput,
|
||||
)
|
||||
|
||||
|
||||
class DiffusionEngine:
|
||||
"Legacy placeholder for DiffusionEngine"
|
||||
|
||||
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")
|
||||
else:
|
||||
config_model_path = getattr(model_cfg, "model", 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 load_model(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def generate(self, *args, **kwargs):
|
||||
"""Generate images based on the input prompt."""
|
||||
# Placeholder implementation
|
||||
# In a full implementation, this would use the loaded model to generate images
|
||||
raise NotImplementedError("This method should be implemented in subclasses.")
|
||||
|
||||
|
||||
class DiffusersPipelineEngine(DiffusionEngine):
|
||||
"""Thin executor that delegates to a diffusion worker."""
|
||||
|
||||
# crucial: must match signature of parent class
|
||||
|
||||
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 or None
|
||||
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", None)
|
||||
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,
|
||||
)
|
||||
@@ -37,7 +37,23 @@ class MultimodalOutputProcessor:
|
||||
handler = self.output_handlers.get(output_type, self._process_pooling_output)
|
||||
|
||||
return handler(engine_core_output)
|
||||
|
||||
|
||||
def _build_request_output(
|
||||
self,
|
||||
source: Any,
|
||||
completion_outputs: List[CompletionOutput],
|
||||
) -> RequestOutput:
|
||||
"""Helper to construct RequestOutput with safe defaults."""
|
||||
|
||||
return RequestOutput(
|
||||
request_id=getattr(source, "request_id", "unknown"),
|
||||
prompt=getattr(source, "prompt", ""),
|
||||
prompt_token_ids=getattr(source, "prompt_token_ids", []),
|
||||
prompt_logprobs=getattr(source, "prompt_logprobs", None),
|
||||
outputs=completion_outputs,
|
||||
finished=getattr(source, "finished", True),
|
||||
)
|
||||
|
||||
def _detect_output_type(self, output: Any) -> str:
|
||||
"""Detect the type of output based on its content."""
|
||||
if hasattr(output, 'output_type'):
|
||||
@@ -76,16 +92,8 @@ class MultimodalOutputProcessor:
|
||||
finish_reason=getattr(output, 'finish_reason', 'length')
|
||||
)
|
||||
|
||||
request_output = RequestOutput(
|
||||
request_id=getattr(output, 'request_id', 'unknown'),
|
||||
prompt=getattr(output, 'prompt', ''),
|
||||
prompt_token_ids=getattr(output, 'prompt_token_ids', []),
|
||||
outputs=[completion_output],
|
||||
finished=getattr(output, 'finished', True)
|
||||
)
|
||||
|
||||
return [request_output]
|
||||
|
||||
return [self._build_request_output(output, [completion_output])]
|
||||
|
||||
def _process_image_output(self, output: Any) -> List[RequestOutput]:
|
||||
"""Process image output."""
|
||||
# For image outputs, we need to create a special RequestOutput
|
||||
@@ -109,15 +117,9 @@ class MultimodalOutputProcessor:
|
||||
# Add image data to the completion output
|
||||
completion_output.image = image_data
|
||||
|
||||
request_output = RequestOutput(
|
||||
request_id=getattr(output, 'request_id', 'unknown'),
|
||||
prompt=getattr(output, 'prompt', ''),
|
||||
prompt_token_ids=getattr(output, 'prompt_token_ids', []),
|
||||
outputs=[completion_output],
|
||||
finished=getattr(output, 'finished', True)
|
||||
)
|
||||
|
||||
return [request_output]
|
||||
return [
|
||||
self._build_request_output(output, [completion_output])
|
||||
]
|
||||
|
||||
def _process_text_image_output(self, output: Any) -> List[RequestOutput]:
|
||||
"""Process combined text and image output."""
|
||||
@@ -141,15 +143,9 @@ class MultimodalOutputProcessor:
|
||||
# Add image data to the completion output
|
||||
completion_output.image = image_data
|
||||
|
||||
request_output = RequestOutput(
|
||||
request_id=getattr(output, 'request_id', 'unknown'),
|
||||
prompt=getattr(output, 'prompt', ''),
|
||||
prompt_token_ids=getattr(output, 'prompt_token_ids', []),
|
||||
outputs=[completion_output],
|
||||
finished=getattr(output, 'finished', True)
|
||||
)
|
||||
|
||||
return [request_output]
|
||||
return [
|
||||
self._build_request_output(output, [completion_output])
|
||||
]
|
||||
|
||||
def _process_latents_output(self, output: Any) -> List[RequestOutput]:
|
||||
"""Process latent representation output."""
|
||||
@@ -171,15 +167,9 @@ class MultimodalOutputProcessor:
|
||||
# Add latent data to the completion output
|
||||
completion_output.latents = latent_data
|
||||
|
||||
request_output = RequestOutput(
|
||||
request_id=getattr(output, 'request_id', 'unknown'),
|
||||
prompt=getattr(output, 'prompt', ''),
|
||||
prompt_token_ids=getattr(output, 'prompt_token_ids', []),
|
||||
outputs=[completion_output],
|
||||
finished=getattr(output, 'finished', True)
|
||||
)
|
||||
|
||||
return [request_output]
|
||||
return [
|
||||
self._build_request_output(output, [completion_output])
|
||||
]
|
||||
|
||||
def _process_pooling_output(self, output: Any) -> List[RequestOutput]:
|
||||
"""Process pooling output (hidden states, embeddings, etc.)."""
|
||||
@@ -201,15 +191,9 @@ class MultimodalOutputProcessor:
|
||||
# Add pooling data to the completion output
|
||||
completion_output.pooler_output = pooling_data
|
||||
|
||||
request_output = RequestOutput(
|
||||
request_id=getattr(output, 'request_id', 'unknown'),
|
||||
prompt=getattr(output, 'prompt', ''),
|
||||
prompt_token_ids=getattr(output, 'prompt_token_ids', []),
|
||||
outputs=[completion_output],
|
||||
finished=getattr(output, 'finished', True)
|
||||
)
|
||||
|
||||
return [request_output]
|
||||
return [
|
||||
self._build_request_output(output, [completion_output])
|
||||
]
|
||||
|
||||
def process_outputs(self, engine_core_outputs: List[EngineCoreOutput], **kwargs) -> List[RequestOutput]:
|
||||
"""Process multiple engine core outputs."""
|
||||
|
||||
@@ -5,8 +5,13 @@ Omni serve command for vLLM-omni.
|
||||
import argparse
|
||||
import asyncio
|
||||
from typing import List, Optional
|
||||
|
||||
from .omni_llm import AsyncOmniLLM
|
||||
from vllm_omni.config import create_ar_stage_config, create_dit_stage_config, DiTConfig
|
||||
from vllm_omni.config import (
|
||||
DiTConfig,
|
||||
create_ar_stage_config,
|
||||
create_dit_stage_config,
|
||||
)
|
||||
|
||||
|
||||
class OmniServeCommand:
|
||||
@@ -129,11 +134,10 @@ class OmniServeCommand:
|
||||
|
||||
if dit_model:
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=args.dit_steps,
|
||||
guidance_scale=args.dit_guidance_scale,
|
||||
use_diffusers=args.use_diffusers
|
||||
use_diffusers=args.use_diffusers,
|
||||
diffusers_pipeline="auto" if args.use_diffusers else None,
|
||||
)
|
||||
|
||||
dit_stage_config = create_dit_stage_config(
|
||||
@@ -160,11 +164,10 @@ class OmniServeCommand:
|
||||
stage_configs.append(ar_config)
|
||||
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=args.dit_steps,
|
||||
guidance_scale=args.dit_guidance_scale,
|
||||
use_diffusers=args.use_diffusers
|
||||
use_diffusers=args.use_diffusers,
|
||||
diffusers_pipeline="auto" if args.use_diffusers else None,
|
||||
)
|
||||
|
||||
dit_stage_config = create_dit_stage_config(
|
||||
|
||||
@@ -8,10 +8,10 @@ from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
from vllm.outputs import RequestOutput, LoRARequest
|
||||
|
||||
|
||||
from ..config import OmniStageConfig
|
||||
from ..core.stage_manager import StageManager
|
||||
from ..engine.output_processor import MultimodalOutputProcessor
|
||||
from ..engine.diffusion_engine import DiffusersPipelineEngine
|
||||
|
||||
|
||||
class OmniLLM(LLM):
|
||||
@@ -25,7 +25,10 @@ class OmniLLM(LLM):
|
||||
):
|
||||
# Use the first stage's model as the default model for LLM
|
||||
default_model = stage_configs[0].model_path if stage_configs else "test-model"
|
||||
|
||||
|
||||
# Track whether we should launch engines in multiprocess mode.
|
||||
self.multiprocess_mode = kwargs.pop("multiprocess_mode", False)
|
||||
|
||||
# Fix configuration validation issues
|
||||
# Ensure max_num_batched_tokens is at least as large as max_model_len
|
||||
if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs:
|
||||
@@ -52,14 +55,24 @@ class OmniLLM(LLM):
|
||||
|
||||
def generate(
|
||||
self,
|
||||
stage_args_list: List[Dict[str, Any]],
|
||||
stage_args_list: Optional[List[Dict[str, Any]]] = None,
|
||||
use_tqdm: Union[bool, Callable[..., Any]] = True,
|
||||
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
||||
priority: Optional[List[int]] = None,
|
||||
**kwargs
|
||||
*,
|
||||
prompt: Optional[str] = None,
|
||||
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
|
||||
) -> List[RequestOutput]:
|
||||
"""Main generation interface - orchestrates multi-stage processing."""
|
||||
|
||||
if stage_args_list is None:
|
||||
if prompt is None:
|
||||
raise ValueError(
|
||||
"prompt must be provided when stage_args_list is not supplied"
|
||||
)
|
||||
stage_args_list = self._build_stage_args_from_config(
|
||||
prompt, stage_overrides or {}
|
||||
)
|
||||
|
||||
if len(stage_args_list) != len(self.stage_configs):
|
||||
raise ValueError(
|
||||
f"Number of stage arguments ({len(stage_args_list)}) must match "
|
||||
@@ -74,12 +87,12 @@ class OmniLLM(LLM):
|
||||
|
||||
# Prepare input for this stage
|
||||
processed_input = self._process_stage_inputs(
|
||||
stage_config, stage_args, current_output
|
||||
stage_config, stage_args or {}, current_output
|
||||
)
|
||||
|
||||
# Execute stage
|
||||
stage_output = self._execute_stage(
|
||||
stage_engine, processed_input, lora_request, priority
|
||||
stage_engine, processed_input, lora_request, priority, stage_config
|
||||
)
|
||||
|
||||
# Update for next stage
|
||||
@@ -90,6 +103,23 @@ class OmniLLM(LLM):
|
||||
final_output = self.output_processor.process_output(current_output)
|
||||
return final_output
|
||||
|
||||
def _build_stage_args_from_config(
|
||||
self,
|
||||
prompt: str,
|
||||
stage_overrides: Dict[int, Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Derive per-stage argument dictionaries from configuration defaults."""
|
||||
stage_args: List[Dict[str, Any]] = []
|
||||
for stage_config in self.stage_configs:
|
||||
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
|
||||
override = stage_overrides.get(stage_config.stage_id)
|
||||
if override:
|
||||
combined.update(override)
|
||||
if stage_config.engine_type == "AR":
|
||||
combined["prompt"] = prompt
|
||||
stage_args.append(combined)
|
||||
return stage_args
|
||||
|
||||
def _process_stage_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
@@ -98,24 +128,24 @@ class OmniLLM(LLM):
|
||||
) -> Dict[str, Any]:
|
||||
"""Prepare input for specific stage."""
|
||||
if stage_config.engine_type == "AR":
|
||||
return self._process_ar_inputs(stage_args, previous_output)
|
||||
return self._process_ar_inputs(stage_config, stage_args, previous_output)
|
||||
elif stage_config.engine_type == "DiT":
|
||||
return self._process_dit_inputs(stage_args, previous_output)
|
||||
return self._process_dit_inputs(stage_config, stage_args, previous_output)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
|
||||
|
||||
|
||||
def _process_ar_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
stage_args: Dict[str, Any],
|
||||
previous_output: Optional[Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Process inputs for AR stage."""
|
||||
# For AR stages, we typically use text prompts
|
||||
processed_input = {
|
||||
"prompt": stage_args.get("prompt", ""),
|
||||
"max_tokens": stage_args.get("max_tokens", 100),
|
||||
"temperature": stage_args.get("temperature", 0.7),
|
||||
}
|
||||
combined = dict(stage_config.default_stage_args or {})
|
||||
combined.update(stage_args)
|
||||
combined.setdefault("prompt", "")
|
||||
combined.setdefault("max_tokens", 100)
|
||||
combined.setdefault("temperature", 0.7)
|
||||
|
||||
# If we have previous output (e.g., from a previous AR stage),
|
||||
# we might want to use it as context
|
||||
@@ -124,49 +154,97 @@ class OmniLLM(LLM):
|
||||
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
||||
last_output = previous_output.outputs[-1]
|
||||
if hasattr(last_output, 'text'):
|
||||
processed_input["prompt"] = last_output.text + " " + processed_input["prompt"]
|
||||
|
||||
return processed_input
|
||||
|
||||
combined["prompt"] = last_output.text + " " + combined["prompt"]
|
||||
|
||||
return combined
|
||||
|
||||
def _process_dit_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
stage_args: Dict[str, Any],
|
||||
previous_output: Optional[Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Process inputs for DiT stage."""
|
||||
processed_input = {
|
||||
"prompt": stage_args.get("prompt", ""),
|
||||
"height": stage_args.get("height", 512),
|
||||
"width": stage_args.get("width", 512),
|
||||
"num_inference_steps": stage_args.get("num_inference_steps", 50),
|
||||
"guidance_scale": stage_args.get("guidance_scale", 7.5),
|
||||
}
|
||||
|
||||
combined = dict(stage_config.default_stage_args or {})
|
||||
combined.update(stage_args)
|
||||
|
||||
dit = stage_config.dit_config
|
||||
if dit is not None:
|
||||
combined.setdefault("height", getattr(dit, "height", 512))
|
||||
combined.setdefault("width", getattr(dit, "width", 512))
|
||||
combined.setdefault(
|
||||
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
|
||||
)
|
||||
combined.setdefault(
|
||||
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
|
||||
)
|
||||
else:
|
||||
combined.setdefault("height", 512)
|
||||
combined.setdefault("width", 512)
|
||||
combined.setdefault("num_inference_steps", 50)
|
||||
combined.setdefault("guidance_scale", 7.5)
|
||||
|
||||
# Handle image inputs if present
|
||||
if "image" in stage_args:
|
||||
# For now, we'll pass the image path directly
|
||||
# In a full implementation, this would involve VAE encoding
|
||||
processed_input["image"] = stage_args["image"]
|
||||
|
||||
combined["image"] = stage_args["image"]
|
||||
|
||||
# If we have previous output from an AR stage, we might want to use it
|
||||
if previous_output is not None:
|
||||
# Extract text from previous AR output
|
||||
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
||||
last_output = previous_output.outputs[-1]
|
||||
if hasattr(last_output, 'text'):
|
||||
processed_input["prompt"] = last_output.text
|
||||
|
||||
return processed_input
|
||||
combined["prompt"] = last_output.text
|
||||
|
||||
combined.setdefault("prompt", stage_args.get("prompt", ""))
|
||||
|
||||
return combined
|
||||
|
||||
def _execute_stage(
|
||||
self,
|
||||
stage_engine: LLMEngine,
|
||||
stage_engine: Optional[Union[LLMEngine, DiffusersPipelineEngine]],
|
||||
processed_input: Dict[str, Any],
|
||||
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
||||
priority: Optional[List[int]] = None
|
||||
priority: Optional[List[int]] = None,
|
||||
stage_config: Optional[OmniStageConfig] = None,
|
||||
) -> Any:
|
||||
"""Execute a single stage."""
|
||||
# Use the parent LLM's generate method for actual text generation
|
||||
# DiT via diffusers backend
|
||||
if stage_config and stage_config.engine_type == "DiT":
|
||||
dit = stage_config.dit_config
|
||||
if dit and getattr(dit, "use_diffusers", False):
|
||||
# Lazy-init executor per stage
|
||||
if not hasattr(self, "_dit_engines"):
|
||||
self._dit_engines = {}
|
||||
exec_inst = self._dit_engines.get(stage_config.stage_id)
|
||||
if exec_inst is None:
|
||||
exec_inst = DiffusersPipelineEngine(
|
||||
dit_config=dit,
|
||||
model_path=stage_config.model_path,
|
||||
log_stats=self.log_stats,
|
||||
multiprocess_mode=self.multiprocess_mode,
|
||||
)
|
||||
|
||||
self._dit_engines[stage_config.stage_id] = exec_inst
|
||||
|
||||
return exec_inst.generate(
|
||||
prompt=processed_input.get("prompt", ""),
|
||||
height=processed_input.get("height", getattr(dit, "height", 512)),
|
||||
width=processed_input.get("width", getattr(dit, "width", 512)),
|
||||
num_inference_steps=processed_input.get(
|
||||
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
|
||||
),
|
||||
guidance_scale=processed_input.get(
|
||||
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
|
||||
),
|
||||
negative_prompt=processed_input.get("negative_prompt"),
|
||||
seed=processed_input.get("seed"),
|
||||
image=processed_input.get("image"),
|
||||
)
|
||||
|
||||
# Use the parent LLM's generate method for AR text generation
|
||||
prompt = processed_input.get("prompt", "")
|
||||
max_tokens = processed_input.get("max_tokens", 100)
|
||||
temperature = processed_input.get("temperature", 0.7)
|
||||
@@ -252,14 +330,25 @@ class AsyncOmniLLM(LLM):
|
||||
|
||||
async def generate_async(
|
||||
self,
|
||||
stage_args_list: List[Dict[str, Any]],
|
||||
stage_args_list: Optional[List[Dict[str, Any]]] = None,
|
||||
use_tqdm: Union[bool, Callable[..., Any]] = True,
|
||||
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
||||
priority: Optional[List[int]] = None,
|
||||
**kwargs
|
||||
*,
|
||||
prompt: Optional[str] = None,
|
||||
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
|
||||
) -> List[RequestOutput]:
|
||||
"""Async generation interface - orchestrates multi-stage processing."""
|
||||
|
||||
|
||||
if stage_args_list is None:
|
||||
if prompt is None:
|
||||
raise ValueError(
|
||||
"prompt must be provided when stage_args_list is not supplied"
|
||||
)
|
||||
stage_args_list = self._build_stage_args_from_config(
|
||||
prompt, stage_overrides or {}
|
||||
)
|
||||
|
||||
if len(stage_args_list) != len(self.stage_configs):
|
||||
raise ValueError(
|
||||
f"Number of stage arguments ({len(stage_args_list)}) must match "
|
||||
@@ -274,12 +363,12 @@ class AsyncOmniLLM(LLM):
|
||||
|
||||
# Prepare input for this stage
|
||||
processed_input = self._process_stage_inputs(
|
||||
stage_config, stage_args, current_output
|
||||
stage_config, stage_args or {}, current_output
|
||||
)
|
||||
|
||||
# Execute stage asynchronously
|
||||
stage_output = await self._execute_stage_async(
|
||||
stage_engine, processed_input, lora_request, priority
|
||||
stage_engine, processed_input, lora_request, priority, stage_config
|
||||
)
|
||||
|
||||
# Update for next stage
|
||||
@@ -290,6 +379,22 @@ class AsyncOmniLLM(LLM):
|
||||
final_output = self.output_processor.process_output(current_output)
|
||||
return final_output
|
||||
|
||||
def _build_stage_args_from_config(
|
||||
self,
|
||||
prompt: str,
|
||||
stage_overrides: Dict[int, Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
stage_args: List[Dict[str, Any]] = []
|
||||
for stage_config in self.stage_configs:
|
||||
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
|
||||
override = stage_overrides.get(stage_config.stage_id)
|
||||
if override:
|
||||
combined.update(override)
|
||||
if stage_config.engine_type == "AR":
|
||||
combined["prompt"] = prompt
|
||||
stage_args.append(combined)
|
||||
return stage_args
|
||||
|
||||
def _process_stage_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
@@ -298,66 +403,133 @@ class AsyncOmniLLM(LLM):
|
||||
) -> Dict[str, Any]:
|
||||
"""Prepare input for specific stage (same as OmniLLM)."""
|
||||
if stage_config.engine_type == "AR":
|
||||
return self._process_ar_inputs(stage_args, previous_output)
|
||||
return self._process_ar_inputs(stage_config, stage_args, previous_output)
|
||||
elif stage_config.engine_type == "DiT":
|
||||
return self._process_dit_inputs(stage_args, previous_output)
|
||||
return self._process_dit_inputs(stage_config, stage_args, previous_output)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
|
||||
|
||||
|
||||
def _process_ar_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
stage_args: Dict[str, Any],
|
||||
previous_output: Optional[Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Process inputs for AR stage (same as OmniLLM)."""
|
||||
processed_input = {
|
||||
"prompt": stage_args.get("prompt", ""),
|
||||
"max_tokens": stage_args.get("max_tokens", 100),
|
||||
"temperature": stage_args.get("temperature", 0.7),
|
||||
}
|
||||
|
||||
combined = dict(stage_config.default_stage_args or {})
|
||||
combined.update(stage_args)
|
||||
combined.setdefault("prompt", "")
|
||||
combined.setdefault("max_tokens", 100)
|
||||
combined.setdefault("temperature", 0.7)
|
||||
|
||||
if previous_output is not None:
|
||||
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
||||
last_output = previous_output.outputs[-1]
|
||||
if hasattr(last_output, 'text'):
|
||||
processed_input["prompt"] = last_output.text + " " + processed_input["prompt"]
|
||||
|
||||
return processed_input
|
||||
|
||||
combined["prompt"] = last_output.text + " " + combined["prompt"]
|
||||
|
||||
return combined
|
||||
|
||||
def _process_dit_inputs(
|
||||
self,
|
||||
stage_config: OmniStageConfig,
|
||||
stage_args: Dict[str, Any],
|
||||
previous_output: Optional[Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Process inputs for DiT stage (same as OmniLLM)."""
|
||||
processed_input = {
|
||||
"prompt": stage_args.get("prompt", ""),
|
||||
"height": stage_args.get("height", 512),
|
||||
"width": stage_args.get("width", 512),
|
||||
"num_inference_steps": stage_args.get("num_inference_steps", 50),
|
||||
"guidance_scale": stage_args.get("guidance_scale", 7.5),
|
||||
}
|
||||
|
||||
combined = dict(stage_config.default_stage_args or {})
|
||||
combined.update(stage_args)
|
||||
|
||||
dit = stage_config.dit_config
|
||||
if dit is not None:
|
||||
combined.setdefault("height", getattr(dit, "height", 512))
|
||||
combined.setdefault("width", getattr(dit, "width", 512))
|
||||
combined.setdefault(
|
||||
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
|
||||
)
|
||||
combined.setdefault(
|
||||
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
|
||||
)
|
||||
else:
|
||||
combined.setdefault("height", 512)
|
||||
combined.setdefault("width", 512)
|
||||
combined.setdefault("num_inference_steps", 50)
|
||||
combined.setdefault("guidance_scale", 7.5)
|
||||
|
||||
if "image" in stage_args:
|
||||
processed_input["image"] = stage_args["image"]
|
||||
|
||||
combined["image"] = stage_args["image"]
|
||||
|
||||
if previous_output is not None:
|
||||
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
||||
last_output = previous_output.outputs[-1]
|
||||
if hasattr(last_output, 'text'):
|
||||
processed_input["prompt"] = last_output.text
|
||||
|
||||
return processed_input
|
||||
combined["prompt"] = last_output.text
|
||||
|
||||
combined.setdefault("prompt", stage_args.get("prompt", ""))
|
||||
|
||||
return combined
|
||||
|
||||
async def _execute_stage_async(
|
||||
self,
|
||||
stage_engine: AsyncLLM,
|
||||
processed_input: Dict[str, Any],
|
||||
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
||||
priority: Optional[List[int]] = None
|
||||
priority: Optional[List[int]] = None,
|
||||
stage_config: Optional[OmniStageConfig] = None,
|
||||
) -> Any:
|
||||
"""Execute a single stage asynchronously."""
|
||||
# Use the parent LLM's generate method for actual text generation
|
||||
# DiT via diffusers backend (sync call inside async for MVP)
|
||||
if stage_config and stage_config.engine_type == "DiT":
|
||||
dit = stage_config.dit_config
|
||||
if dit and getattr(dit, "use_diffusers", False):
|
||||
if not hasattr(self, "_dit_executors"):
|
||||
self._dit_executors = {}
|
||||
exec_inst = self._dit_executors.get(stage_config.stage_id)
|
||||
if exec_inst is None:
|
||||
from vllm_omni.executor.diffusers_executor import (
|
||||
DiffusersPipelineExecutor,
|
||||
)
|
||||
|
||||
pipeline_name = getattr(dit, "diffusers_pipeline", None)
|
||||
device_cfg = getattr(dit, "device_config", None)
|
||||
model_cfg = getattr(dit, "model_config", None)
|
||||
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:
|
||||
if isinstance(model_cfg, dict):
|
||||
dtype = model_cfg.get("dtype")
|
||||
else:
|
||||
dtype = getattr(model_cfg, "dtype", None)
|
||||
|
||||
exec_inst = DiffusersPipelineExecutor(
|
||||
model_path=stage_config.model_path,
|
||||
pipeline_name=pipeline_name,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
self._dit_executors[stage_config.stage_id] = exec_inst
|
||||
|
||||
return exec_inst.generate(
|
||||
prompt=processed_input.get("prompt", ""),
|
||||
height=processed_input.get("height", getattr(dit, "height", 512)),
|
||||
width=processed_input.get("width", getattr(dit, "width", 512)),
|
||||
num_inference_steps=processed_input.get(
|
||||
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
|
||||
),
|
||||
guidance_scale=processed_input.get(
|
||||
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
|
||||
),
|
||||
negative_prompt=processed_input.get("negative_prompt"),
|
||||
seed=processed_input.get("seed"),
|
||||
image=processed_input.get("image"),
|
||||
)
|
||||
|
||||
# Use the parent LLM's generate method for AR text generation
|
||||
prompt = processed_input.get("prompt", "")
|
||||
max_tokens = processed_input.get("max_tokens", 100)
|
||||
temperature = processed_input.get("temperature", 0.7)
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Diffusers-backed executor for DiT stages.
|
||||
|
||||
Provides a minimal wrapper around HuggingFace diffusers pipelines to support
|
||||
text-to-image generation as the DiT stage backend.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from vllm_omni.worker.gpu_diffusion_worker import (
|
||||
DiffusionGPUWorker,
|
||||
DiffusionRunnerOutput,
|
||||
)
|
||||
|
||||
|
||||
class DiffusersPipelineExecutor:
|
||||
"""Thin executor that delegates to a diffusion worker."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
device: Optional[str] = None,
|
||||
dtype: Optional[str] = None,
|
||||
pipeline_name: Optional[str] = None,
|
||||
) -> None:
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
import diffusers # noqa: F401
|
||||
except Exception as e: # pragma: no cover
|
||||
raise ImportError(
|
||||
"diffusers and torch are required for DiT diffusers backend."
|
||||
) from e
|
||||
|
||||
self.worker = DiffusionGPUWorker(
|
||||
model_path,
|
||||
pipeline_name=pipeline_name,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 5.0,
|
||||
negative_prompt: Optional[str] = None,
|
||||
seed: Optional[int] = None,
|
||||
image: Optional[Any] = None,
|
||||
) -> DiffusionRunnerOutput:
|
||||
return self.worker.generate(
|
||||
prompt=prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
image=image,
|
||||
)
|
||||
Whitespace-only changes.
@@ -1,166 +0,0 @@
|
||||
"""
|
||||
vLLM plugin system for vLLM-omni integration.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional
|
||||
from vllm_omni.entrypoints.omni_llm import OmniLLM, AsyncOmniLLM
|
||||
from vllm_omni.config import create_ar_stage_config, create_dit_stage_config
|
||||
|
||||
|
||||
class OmniPlugin:
|
||||
"""vLLM plugin for vLLM-omni integration."""
|
||||
|
||||
def __init__(self):
|
||||
self.name = "omni"
|
||||
self.version = "0.1.0"
|
||||
self.description = "Multi-modality models inference and serving"
|
||||
|
||||
def register_components(self) -> Dict[str, Any]:
|
||||
"""Register vLLM-omni components with vLLM."""
|
||||
return {
|
||||
"omni_llm": OmniLLM,
|
||||
"async_omni_llm": AsyncOmniLLM,
|
||||
"create_ar_stage_config": create_ar_stage_config,
|
||||
"create_dit_stage_config": create_dit_stage_config,
|
||||
}
|
||||
|
||||
def get_config_schema(self) -> Dict[str, Any]:
|
||||
"""Get configuration schema for the plugin."""
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"stages": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"stage_id": {"type": "integer"},
|
||||
"engine_type": {"type": "string", "enum": ["AR", "DiT"]},
|
||||
"model_path": {"type": "string"},
|
||||
"input_modalities": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"}
|
||||
},
|
||||
"output_modalities": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"}
|
||||
}
|
||||
},
|
||||
"required": ["stage_id", "engine_type", "model_path", "input_modalities", "output_modalities"]
|
||||
}
|
||||
}
|
||||
},
|
||||
"required": ["stages"]
|
||||
}
|
||||
|
||||
def validate_config(self, config: Dict[str, Any]) -> bool:
|
||||
"""Validate plugin configuration."""
|
||||
if "stages" not in config:
|
||||
return False
|
||||
|
||||
stages = config["stages"]
|
||||
if not isinstance(stages, list):
|
||||
return False
|
||||
|
||||
for stage in stages:
|
||||
required_fields = ["stage_id", "engine_type", "model_path", "input_modalities", "output_modalities"]
|
||||
for field in required_fields:
|
||||
if field not in stage:
|
||||
return False
|
||||
|
||||
if stage["engine_type"] not in ["AR", "DiT"]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def create_omni_llm(self, config: Dict[str, Any]) -> OmniLLM:
|
||||
"""Create an OmniLLM instance from configuration."""
|
||||
from vllm_omni.config import OmniStageConfig, DiTConfig
|
||||
|
||||
stage_configs = []
|
||||
for stage_config in config["stages"]:
|
||||
if stage_config["engine_type"] == "AR":
|
||||
stage_config_obj = create_ar_stage_config(
|
||||
stage_id=stage_config["stage_id"],
|
||||
model_path=stage_config["model_path"],
|
||||
input_modalities=stage_config["input_modalities"],
|
||||
output_modalities=stage_config["output_modalities"]
|
||||
)
|
||||
elif stage_config["engine_type"] == "DiT":
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=50,
|
||||
guidance_scale=7.5
|
||||
)
|
||||
|
||||
stage_config_obj = create_dit_stage_config(
|
||||
stage_id=stage_config["stage_id"],
|
||||
model_path=stage_config["model_path"],
|
||||
input_modalities=stage_config["input_modalities"],
|
||||
output_modalities=stage_config["output_modalities"],
|
||||
dit_config=dit_config
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown engine type: {stage_config['engine_type']}")
|
||||
|
||||
stage_configs.append(stage_config_obj)
|
||||
|
||||
return OmniLLM(stage_configs)
|
||||
|
||||
def create_async_omni_llm(self, config: Dict[str, Any]) -> AsyncOmniLLM:
|
||||
"""Create an AsyncOmniLLM instance from configuration."""
|
||||
from vllm_omni.config import OmniStageConfig, DiTConfig
|
||||
|
||||
stage_configs = []
|
||||
for stage_config in config["stages"]:
|
||||
if stage_config["engine_type"] == "AR":
|
||||
stage_config_obj = create_ar_stage_config(
|
||||
stage_id=stage_config["stage_id"],
|
||||
model_path=stage_config["model_path"],
|
||||
input_modalities=stage_config["input_modalities"],
|
||||
output_modalities=stage_config["output_modalities"]
|
||||
)
|
||||
elif stage_config["engine_type"] == "DiT":
|
||||
dit_config = DiTConfig(
|
||||
model_type="dit",
|
||||
scheduler_type="ddpm",
|
||||
num_inference_steps=50,
|
||||
guidance_scale=7.5
|
||||
)
|
||||
|
||||
stage_config_obj = create_dit_stage_config(
|
||||
stage_id=stage_config["stage_id"],
|
||||
model_path=stage_config["model_path"],
|
||||
input_modalities=stage_config["input_modalities"],
|
||||
output_modalities=stage_config["output_modalities"],
|
||||
dit_config=dit_config
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown engine type: {stage_config['engine_type']}")
|
||||
|
||||
stage_configs.append(stage_config_obj)
|
||||
|
||||
return AsyncOmniLLM(stage_configs)
|
||||
|
||||
def get_help_text(self) -> str:
|
||||
"""Get help text for the plugin."""
|
||||
return """
|
||||
vLLM-omni Plugin
|
||||
|
||||
This plugin enables multi-modality models inference and serving with non-autoregressive structures.
|
||||
|
||||
Usage:
|
||||
vllm serve model --omni [options]
|
||||
|
||||
Options:
|
||||
--ar-stage MODEL_PATH AR stage model path
|
||||
--dit-stage MODEL_PATH DiT stage model path
|
||||
--dit-steps N Number of DiT inference steps (default: 50)
|
||||
--dit-guidance-scale F DiT guidance scale (default: 7.5)
|
||||
--use-diffusers Use diffusers pipeline for DiT stage
|
||||
|
||||
Examples:
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni
|
||||
vllm serve model --omni --ar-stage text-model --dit-stage image-model
|
||||
"""
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Model runner that wraps a diffusers text-to-image pipeline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
@dataclass
|
||||
class DiffusionRunnerOutput:
|
||||
"""Structured output returned by the diffusion runner."""
|
||||
|
||||
prompt: str
|
||||
images: list
|
||||
request_id: str = "diffusion_request"
|
||||
finished: bool = True
|
||||
output_type: str = "image"
|
||||
|
||||
|
||||
class DiffusionModelRunner:
|
||||
"""Lightweight runner that loads and executes diffusers pipelines."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
*,
|
||||
pipeline_name: Optional[str] = None,
|
||||
device: Optional[str] = None,
|
||||
dtype: Optional[str] = None,
|
||||
) -> None:
|
||||
self.model_path = model_path
|
||||
self.pipeline_name = pipeline_name or "auto"
|
||||
self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype)
|
||||
self._pipeline = self._load_pipeline()
|
||||
|
||||
def _resolve_device_and_dtype(
|
||||
self,
|
||||
device: Optional[str],
|
||||
dtype: Optional[str],
|
||||
):
|
||||
import torch
|
||||
|
||||
resolved_device = device
|
||||
if resolved_device is None:
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
resolved_device = "mps"
|
||||
elif torch.cuda.is_available():
|
||||
resolved_device = "cuda"
|
||||
else:
|
||||
resolved_device = "cpu"
|
||||
|
||||
if dtype is not None:
|
||||
resolved_dtype = getattr(torch, dtype, None)
|
||||
else:
|
||||
if resolved_device in {"cuda", "mps"}:
|
||||
resolved_dtype = torch.float16
|
||||
else:
|
||||
resolved_dtype = torch.float32
|
||||
|
||||
return resolved_device, resolved_dtype
|
||||
|
||||
def _load_pipeline(self):
|
||||
import diffusers
|
||||
|
||||
if self.pipeline_name != "auto":
|
||||
PipelineCls = getattr(diffusers, self.pipeline_name)
|
||||
else:
|
||||
PipelineCls = diffusers.AutoPipelineForText2Image
|
||||
|
||||
pipeline = PipelineCls.from_pretrained(
|
||||
self.model_path,
|
||||
torch_dtype=self.torch_dtype,
|
||||
)
|
||||
|
||||
try:
|
||||
pipeline = pipeline.to(self.device)
|
||||
except Exception:
|
||||
pipeline = pipeline.to("cpu")
|
||||
|
||||
self._apply_optimisations(pipeline)
|
||||
return pipeline
|
||||
|
||||
@staticmethod
|
||||
def _apply_optimisations(pipeline) -> None:
|
||||
# Enable attention slicing / VAE tiling where available to reduce memory.
|
||||
try:
|
||||
if hasattr(pipeline, "enable_attention_slicing"):
|
||||
pipeline.enable_attention_slicing()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
vae = getattr(pipeline, "vae", None)
|
||||
if vae and hasattr(vae, "enable_tiling"):
|
||||
vae.enable_tiling()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 5.0,
|
||||
negative_prompt: Optional[str] = None,
|
||||
seed: Optional[int] = None,
|
||||
image: Optional[Any] = None,
|
||||
) -> DiffusionRunnerOutput:
|
||||
import torch
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
try:
|
||||
generator = torch.Generator(device=self.device).manual_seed(int(seed))
|
||||
except Exception:
|
||||
# Some backends (e.g., CPU without full support) may not accept the device.
|
||||
generator = torch.Generator().manual_seed(int(seed))
|
||||
|
||||
output = self._pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
num_inference_steps=int(num_inference_steps),
|
||||
guidance_scale=float(guidance_scale),
|
||||
height=int(height),
|
||||
width=int(width),
|
||||
generator=generator,
|
||||
image=image,
|
||||
)
|
||||
|
||||
return DiffusionRunnerOutput(
|
||||
prompt=prompt,
|
||||
images=getattr(output, "images", []) or [],
|
||||
)
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Worker abstraction for executing diffusion model runners."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from vllm.v1.worker.worker_base import WorkerBase
|
||||
|
||||
from .gpu_diffusion_model_runner import DiffusionModelRunner, DiffusionRunnerOutput
|
||||
|
||||
|
||||
class DiffusionGPUWorker(WorkerBase):
|
||||
"""Minimal worker wrapping DiffusionModelRunner.
|
||||
|
||||
The worker interface mirrors the shape expected by vLLM executors: it owns
|
||||
a model runner instance and exposes a `generate` method that can be called
|
||||
by an executor implementation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
*,
|
||||
pipeline_name: Optional[str] = None,
|
||||
device: Optional[str] = None,
|
||||
dtype: Optional[str] = None,
|
||||
) -> None:
|
||||
self.model_runner = DiffusionModelRunner(
|
||||
model_path,
|
||||
pipeline_name=pipeline_name,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 5.0,
|
||||
negative_prompt: Optional[str] = None,
|
||||
seed: Optional[int] = None,
|
||||
image: Optional[Any] = None,
|
||||
) -> DiffusionRunnerOutput:
|
||||
return self.model_runner.generate(
|
||||
prompt=prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
image=image,
|
||||
)
|
||||
Reference in new issue
Block a user