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

[Misc] Support WorkerWrapperBase and CustomPipeline for Diffusion Worker (#764)

Browse Source 
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
This commit is contained in:
knlnguyen1802 authored and GitHub committed 2026-02-12 19:23:51 +08:00
1 parent e427da022b
commit 9994594782
12 files changed
+1359 -23
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/features/custom_pipeline.md
+151
View File
@@ -0,0 +1,151 @@
# Custom Pipeline Extension Guide
Transformer already support Custom Pipeline via
https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md
This guide demonstrates how to use the newly added features for extending vLLM-Omni's diffusion pipeline with custom functionality.
## Overview
Three main features enable custom pipeline extension:
1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
## Features
### WorkerWrapperBase
`WorkerWrapperBase` is a wrapper class that creates `DiffusionWorker` instances with optional extension support. It enables dynamic inheritance, allowing you to add custom methods and functionality to workers without modifying the base worker class.
**Key capabilities:**
- Dynamic worker class extension via `worker_extension_cls`
- Support for custom pipeline initialization via `custom_pipeline_args`
- Method delegation to underlying worker
- Attribute access forwarding
**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
### load_format Parameter
The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
- **`"default"`**: Standard model loading using the model registry (default behavior)
- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
### CustomPipelineWorkerExtension
`CustomPipelineWorkerExtension` is a mixin class that extends `DiffusionWorker` with the ability to re-initialize the pipeline with a custom implementation.
**Key method:**
- `re_init_pipeline(custom_pipeline_args)`: Re-initializes the pipeline with custom arguments, properly cleaning up the old pipeline first
**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
## Usage Example
### Step 1: Create a Custom Pipeline
Create a custom pipeline class that extends an existing pipeline. In this example, we extend `QwenImageEditPipeline` to add trajectory tracking:
```python
# custom_pipeline.py
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
import torch
class CustomPipeline(QwenImageEditPipeline):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
def forward(self, req, prompt=None, negative_prompt=None, **kwargs):
# Call parent's forward to get normal output
output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs)
# Add custom trajectory data
actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50)
output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32)
output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
return output
```
### Step 2: Use the Custom Pipeline with Omni
Initialize the `Omni` engine with custom pipeline configuration:
```python
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Initialize with custom pipeline
omni = Omni(
model="Qwen/Qwen-Image-Edit",
diffusion_load_format="dummy", # Skip initial loading
custom_pipeline_args={
"pipeline_class": "custom_pipeline.CustomPipeline"
},
)
# Generate with the custom pipeline
outputs = omni.generate(
...
)
# Access custom trajectory data
output = outputs[0].request_output[0]
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
print(f"Trajectory latents shape: {output.latents.shape}")
```
### Step 3: Run the Example
The example provided in this directory demonstrates the complete workflow:
```bash
cd examples/offline_inference/custom_pipeline/image_to_image
# Run with custom pipeline
python image_edit.py \
--model Qwen/Qwen-Image-Edit-2511 \
--image cherry_blossom.jpg \
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
--output output_image_edit.png \
--num_inference_steps 10
```
## Advanced Usage
### Custom Worker Extension
You can create custom worker extensions to add new methods beyond pipeline re-initialization:
```python
from typing import Any
from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
class MyCustomExtension(CustomPipelineWorkerExtension):
def custom_method(self):
"""Your custom worker method."""
return "custom_result"
def another_method(self, data: Any):
"""Another custom method."""
# Access worker internals via self
return self.model_runner.some_operation(data)
omni = Omni(
model="Qwen/Qwen-Image-Edit",
diffusion_load_format="dummy",
custom_pipeline_args={
"pipeline_class": "custom_pipeline.CustomPipeline"
},
worker_extension_cls=MyCustomExtension,
# Note: worker_extension_cls is an internal parameter
# CustomPipelineWorkerExtension will automatically init pipeline when custom_pipeline_args is provided
)
```
examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py
+80
View File
@@ -0,0 +1,80 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
from typing import Any
import PIL.Image
import torch
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
logger = logging.getLogger(__name__)
class CustomPipeline(QwenImageEditPipeline):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
def forward(
self,
req: OmniDiffusionRequest,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
image: PIL.Image.Image | torch.Tensor | None = None,
true_cfg_scale: float = 4.0,
height: int | None = None,
width: int | None = None,
num_inference_steps: int = 50,
sigmas: list[float] | None = None,
guidance_scale: float = 1.0,
num_images_per_prompt: int = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
prompt_embeds_mask: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds_mask: torch.Tensor | None = None,
output_type: str | None = "pil",
attention_kwargs: dict[str, Any] | None = None,
callback_on_step_end_tensor_inputs: list[str] = ["latents"],
max_sequence_length: int = 512,
) -> DiffusionOutput:
"""Forward pass for image editing with dummy trajectory data."""
# Call parent's forward to get the normal output
output = super().forward(
req=req,
prompt=prompt,
negative_prompt=negative_prompt,
image=image,
true_cfg_scale=true_cfg_scale,
height=height,
width=width,
num_inference_steps=num_inference_steps,
sigmas=sigmas,
guidance_scale=guidance_scale,
num_images_per_prompt=num_images_per_prompt,
generator=generator,
latents=latents,
prompt_embeds=prompt_embeds,
prompt_embeds_mask=prompt_embeds_mask,
negative_prompt_embeds=negative_prompt_embeds,
negative_prompt_embeds_mask=negative_prompt_embeds_mask,
output_type=output_type,
attention_kwargs=attention_kwargs,
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
max_sequence_length=max_sequence_length,
)
# Get actual num_inference_steps used
actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps
# Create dummy trajectory data
dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
# Inject dummy trajectory data into output
output.trajectory_latents = dummy_trajectory_latents
return output
examples/offline_inference/custom_pipeline/image_to_image/image_edit.py
+262
View File
@@ -0,0 +1,262 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
======================================================
Qwen Image Editor (via vLLM-Omni)
======================================================
Example CLI for Qwen-Image-Edit and compatible models.
This script edits, combines, or layers images according to a text prompt
using the vLLM Omni diffusion backend.
Examples
---------
Single image edit:
python image_edit.py \
--image input.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles saying 'Be Kind'" \
--output output.png
Multiple image composition:
python image_edit.py \
--image input1.png input2.png \
--prompt "Combine these into a single magical composition" \
--output combined.png
Accelerated with cache_dit:
python image_edit.py \
--image input.png \
--prompt "Edit description" \
--cache_backend cache_dit \
--cache_dit_enable_taylorseer
Layered RGBA output:
python image_edit.py \
--model "Qwen/Qwen-Image-Layered" \
--image input.png \
--prompt "" \
--output layered \
--layers 4 \
--color-format RGBA
"""
import argparse
import asyncio
import os
import time
from pathlib import Path
import torch
from PIL import Image
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
# ===========================
# Argument Parser
# ===========================
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Edit or generate images using Qwen-Image-Edit.")
parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
parser.add_argument("--negative_prompt", type=str, default=None)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--cfg_scale", type=float, default=4.0)
parser.add_argument("--guidance_scale", type=float, default=1.0)
parser.add_argument("--output", type=str, default="output.png")
parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
parser.add_argument("--ulysses_degree", type=int, default=1)
parser.add_argument("--ring_degree", type=int, default=1)
parser.add_argument("--tensor_parallel_size", type=int, default=1)
parser.add_argument("--layers", type=int, default=4)
parser.add_argument("--resolution", type=int, default=640)
parser.add_argument("--color-format", type=str, default="RGB")
# Acceleration + Optimization Options
parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
parser.add_argument(
"--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
)
parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
parser.add_argument("--enforce_eager", action="store_true")
parser.add_argument("--vae_use_slicing", action="store_true")
parser.add_argument("--vae_use_tiling", action="store_true")
parser.add_argument("--enable-cpu-offload", action="store_true")
return parser.parse_args()
# ===========================
# Core Logic
# ===========================
async def main():
args = parse_args()
# ---- Load input images ----
input_images: list[Image.Image] = []
for image_path in args.image:
if not os.path.exists(image_path):
raise FileNotFoundError(f"Input image not found: {image_path}")
input_images.append(Image.open(image_path).convert(args.color_format))
# Single or multi-image input
input_image: Image.Image | list[Image.Image] = input_images[0] if len(input_images) == 1 else input_images
# ---- Torch setup ----
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
)
# ---- Cache Config ----
if args.cache_backend == "cache_dit":
cache_config = {
"Fn_compute_blocks": args.cache_dit_fn_compute_blocks,
"Bn_compute_blocks": args.cache_dit_bn_compute_blocks,
"max_warmup_steps": args.cache_dit_max_warmup_steps,
"residual_diff_threshold": args.cache_dit_residual_diff_threshold,
"max_continuous_cached_steps": args.cache_dit_max_continuous_cached_steps,
"enable_taylorseer": args.cache_dit_enable_taylorseer,
"taylorseer_order": args.cache_dit_taylorseer_order,
"scm_steps_mask_policy": args.cache_dit_scm_steps_mask_policy,
"scm_steps_policy": args.cache_dit_scm_steps_policy,
}
elif args.cache_backend == "tea_cache":
cache_config = {"rel_l1_thresh": args.tea_cache_rel_l1_thresh}
else:
cache_config = None
# ---- Initialize Omni ----
omni = Omni(
model=args.model,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
diffusion_load_format="dummy",
custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
)
print(">>> Pipeline loaded successfully")
# ---- Profiling + Info ----
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
print(f"\n{'=' * 60}")
print("Generation Configuration")
print(f"Model: {args.model}")
print(f"Inference steps: {args.num_inference_steps}")
print(f"Cache: {args.cache_backend or 'None'}")
print(f"Input(s): {len(input_images)} image(s)")
for idx, img in enumerate(input_images):
print(f" • Image {idx + 1}: size={img.size}")
print(
f"Parallel config: Ulysses={args.ulysses_degree}, Ring={args.ring_degree}, "
f"CFG Parallel={args.cfg_parallel_size}, Tensor Parallel={args.tensor_parallel_size}"
)
print(f"{'=' * 60}\n")
if profiler_enabled:
omni.start_profile()
# ---- Generation ----
t0 = time.perf_counter()
outputs = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"image": input_image},
},
OmniDiffusionSamplingParams(
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
resolution=args.resolution,
),
)
t1 = time.perf_counter()
print(f"Generation completed in {t1 - t0:.2f}s")
if profiler_enabled:
omni.stop_profile()
# ---- Output Validation ----
if not outputs:
raise ValueError("No output produced from omni.generate()")
first_out = outputs[0].request_output[0]
req_out: OmniRequestOutput = first_out
# Verify trajectory data (from custom pipeline)
print(f"\n{'=' * 60}")
print("TRAJECTORY VERIFICATION:")
assert hasattr(req_out, "latents") and req_out.latents is not None
print(f" ✓ trajectory_latents shape: {req_out.latents.shape}")
print(f" ✓ trajectory_latents dtype: {req_out.latents.dtype}")
print(f" ✓ trajectory_latents mean: {req_out.latents.mean().item():.6f}")
print(f" ✓ trajectory_latents std: {req_out.latents.std().item():.6f}")
assert req_out.latents.shape[0] == args.num_inference_steps, (
f"Expected {args.num_inference_steps} latent snapshots, got {req_out.latents.shape[0]}"
)
print(f"{'=' * 60}\n")
# ---- Save images ----
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
suffix = output_path.suffix or ".png"
stem = output_path.stem or "output"
images = req_out.images
if args.num_outputs_per_prompt <= 1:
img = images[0]
if isinstance(img, list): # Layered
for i, sub_img in enumerate(img):
save_path = output_path.parent / f"{stem}_{i}{suffix}"
sub_img.save(save_path)
print(f"Saved layer {i}: {save_path.resolve()}")
else:
img.save(output_path)
print(f"Saved: {output_path.resolve()}")
else:
for idx, img in enumerate(images):
if isinstance(img, list):
for sub_idx, sub_img in enumerate(img):
save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}"
sub_img.save(save_path)
print(f"Saved composite: {save_path.resolve()}")
else:
save_path = output_path.parent / f"{stem}_{idx}{suffix}"
img.save(save_path)
print(f"Saved: {save_path.resolve()}")
# ===========================
# Entrypoint
# ===========================
if __name__ == "__main__":
asyncio.run(main())
examples/offline_inference/custom_pipeline/image_to_image/run.sh
+6
View File
@@ -0,0 +1,6 @@
python image_edit.py \
--model Qwen/Qwen-Image-Edit-2511 \
--image cherry_blossom.jpg \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 10
tests/diffusion/test_diffusion_worker.py
+1 -1
View File
@@ -36,7 +36,7 @@ def mock_od_config():
@pytest.fixture
def mock_gpu_worker(mock_od_config):
"""Create a DiffusionWorker with mocked initialization."""
with patch.object(DiffusionWorker, "init_device"):
with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"):
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
# Mock the model_runner with pipeline
worker.model_runner = Mock()
tests/diffusion/test_worker_wrapper_base.py
+540
View File
@@ -0,0 +1,540 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for WorkerWrapperBase class.
This module tests the WorkerWrapperBase implementation:
- Initialization with and without worker extensions
- Custom pipeline initialization
- Method delegation via execute_method
- Attribute delegation via __getattr__
- Dynamic worker class extension
"""
from typing import Any
from unittest.mock import Mock, patch
import pytest
from vllm_omni.diffusion.worker.diffusion_worker import (
CustomPipelineWorkerExtension,
DiffusionWorker,
WorkerWrapperBase,
)
# -------------------------------------------------------------------------
# Fixtures
# -------------------------------------------------------------------------
@pytest.fixture
def mock_od_config():
"""Create a mock OmniDiffusionConfig for use in tests."""
config = Mock()
config.num_gpus = 1
config.master_port = 12345
config.enable_sleep_mode = False
config.cache_backend = None
config.cache_config = None
config.model = "test-model"
config.diffusion_load_format = None
config.dtype = "float32"
config.max_cpu_loras = 0
config.lora_path = None
config.lora_scale = 1.0
return config
class TestExtension:
"""Simple test extension adding one custom method."""
def custom_method(self):
return "extension_method"
class MockCustomPipeline:
"""Mock custom pipeline for testing."""
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
def __call__(self, *args, **kwargs):
return "pipeline_output"
# -------------------------------------------------------------------------
# Tests: Initialization
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseInitialization:
"""Test WorkerWrapperBase initialization behavior."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_basic_initialization(self, mock_worker_init, mock_od_config):
"""Test basic initialization without extensions."""
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
)
assert wrapper.gpu_id == 0
assert wrapper.od_config == mock_od_config
assert wrapper.base_worker_class == DiffusionWorker
assert wrapper.worker_extension_cls is None
assert wrapper.custom_pipeline_args is None
assert wrapper.worker is not None
mock_worker_init.assert_called_once_with(
local_rank=0,
rank=0,
od_config=mock_od_config,
)
# -------------------------------------------------------------------------
# Tests: Worker Extension Functionality
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseExtension:
"""Test WorkerWrapperBase worker extension functionality."""
def test_prepare_worker_class_without_extension(self, mock_od_config):
"""Test _prepare_worker_class without a worker extension."""
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
)
worker_class = wrapper._prepare_worker_class()
assert worker_class == DiffusionWorker
def test_prepare_worker_class_with_extension_class(self, mock_od_config):
"""Test _prepare_worker_class with an explicit extension class."""
class TestExtension:
def custom_method(self):
return "extension_method"
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
assert hasattr(wrapper.worker.__class__, "custom_method")
assert TestExtension in wrapper.worker.__class__.__bases__
@patch("vllm.utils.import_utils.resolve_obj_by_qualname")
def test_prepare_worker_class_with_extension_string(self, mock_resolve, mock_od_config):
"""Test _prepare_worker_class with worker extension as string."""
mock_resolve.return_value = TestExtension
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension",
)
assert hasattr(wrapper.worker.__class__, "custom_method")
# -------------------------------------------------------------------------
# Tests: Method Delegation
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseDelegation:
"""Test WorkerWrapperBase delegation to wrapped worker."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_generate_delegation(self, mock_worker_init, mock_od_config):
"""Test that generate() delegates to worker.generate()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
mock_output = Mock()
wrapper.worker.generate = Mock(return_value=mock_output)
mock_requests = [Mock()]
result = wrapper.generate(mock_requests)
wrapper.worker.generate.assert_called_once_with(mock_requests)
assert result == mock_output
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_model_delegation(self, mock_worker_init, mock_od_config):
"""Test that execute_model() delegates to worker.execute_model()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
mock_output = Mock()
wrapper.worker.execute_model = Mock(return_value=mock_output)
mock_reqs = [Mock()]
result = wrapper.execute_model(mock_reqs, mock_od_config)
wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config)
assert result == mock_output
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_load_weights_delegation(self, mock_worker_init, mock_od_config):
"""Test that load_weights() delegates to worker.load_weights()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
expected_result = {"weight1", "weight2"}
wrapper.worker.load_weights = Mock(return_value=expected_result)
mock_weights = [("weight1", Mock()), ("weight2", Mock())]
result = wrapper.load_weights(mock_weights)
wrapper.worker.load_weights.assert_called_once_with(mock_weights)
assert result == expected_result
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_sleep_delegation(self, mock_worker_init, mock_od_config):
"""Test that sleep() delegates to worker.sleep()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.sleep = Mock(return_value=True)
result = wrapper.sleep(level=1)
wrapper.worker.sleep.assert_called_once_with(1)
assert result is True
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_wake_up_delegation(self, mock_worker_init, mock_od_config):
"""Test that wake_up() delegates to worker.wake_up()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.wake_up = Mock(return_value=True)
result = wrapper.wake_up(tags=["weights"])
wrapper.worker.wake_up.assert_called_once_with(["weights"])
assert result is True
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_shutdown_delegation(self, mock_worker_init, mock_od_config):
"""Test that shutdown() delegates to worker.shutdown()."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.shutdown = Mock(return_value=None)
result = wrapper.shutdown()
wrapper.worker.shutdown.assert_called_once()
assert result is None
# -------------------------------------------------------------------------
# Tests: execute_method
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseExecuteMethod:
"""Test WorkerWrapperBase.execute_method functionality."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_success(self, mock_worker_init, mock_od_config):
"""Test execute_method successfully calls worker method."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.test_method = Mock(return_value="method_result")
result = wrapper.execute_method("test_method", "arg1", kwarg1="value1")
wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1")
assert result == "method_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_with_no_args(self, mock_worker_init, mock_od_config):
"""Test execute_method with no arguments."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.no_args_method = Mock(return_value="no_args_result")
result = wrapper.execute_method("no_args_method")
wrapper.worker.no_args_method.assert_called_once_with()
assert result == "no_args_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_error(self, mock_worker_init, mock_od_config):
"""Test execute_method raises exception on error."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error"))
with pytest.raises(RuntimeError, match="Test error"):
wrapper.execute_method("error_method")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_invalid_type(self, mock_worker_init, mock_od_config):
"""Test execute_method with invalid method type."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
with pytest.raises(AssertionError, match="Method must be str"):
wrapper.execute_method(b"bytes_method")
# -------------------------------------------------------------------------
# Tests: __getattr__ delegation
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseGetAttr:
"""Test WorkerWrapperBase.__getattr__ delegation."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_delegation(self, mock_worker_init, mock_od_config):
"""Test __getattr__ delegates to worker attributes."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.custom_attribute = "test_value"
assert wrapper.custom_attribute == "test_value"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_method_access(self, mock_worker_init, mock_od_config):
"""Test __getattr__ delegates to worker methods."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.custom_method = Mock(return_value="method_result")
result = wrapper.custom_method()
wrapper.worker.custom_method.assert_called_once()
assert result == "method_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_missing_attribute(self, mock_worker_init, mock_od_config):
"""Test __getattr__ raises AttributeError for missing attributes."""
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
with pytest.raises(AttributeError):
_ = wrapper.nonexistent_attribute
# -------------------------------------------------------------------------
# Tests: Edge Cases
# -------------------------------------------------------------------------
class TestWorkerWrapperBaseEdgeCases:
"""Test WorkerWrapperBase edge cases and special scenarios."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_extension_conflict_warning(self, mock_worker_init, mock_od_config, caplog):
"""Test a warning is logged when an extension conflicts with worker."""
class ConflictExtension:
def load_model(self):
return "extension_load_model"
with patch.object(DiffusionWorker, "load_model"):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=ConflictExtension,
)
assert wrapper.worker is not None
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_multiple_extensions_same_class(self, mock_worker_init, mock_od_config):
"""Test that applying same extension twice doesn't duplicate it."""
class TestExtension:
def custom_method(self):
return "extension"
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper1 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
wrapper2 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
assert hasattr(wrapper1.worker, "custom_method")
assert hasattr(wrapper2.worker, "custom_method")
# -------------------------------------------------------------------------
# Tests: CustomPipelineWorkerExtension
# -------------------------------------------------------------------------
class TestCustomPipelineWorkerExtension:
"""Test CustomPipelineWorkerExtension functionality."""
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_basic(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
"""Test basic re_init_pipeline functionality."""
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomPipelineWorkerExtension,
)
# Setup mock model_runner and pipeline
mock_model_runner = Mock()
mock_pipeline = Mock()
mock_model_runner.pipeline = mock_pipeline
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
# Call re_init_pipeline
wrapper.worker.re_init_pipeline(custom_args)
# Verify load_model was called with correct arguments
wrapper.worker.load_model.assert_called_once_with(
load_format="custom_pipeline",
custom_pipeline_name="tests.diffusion.test_worker_wrapper_base.MockCustomPipeline",
)
wrapper.worker.init_lora_manager.assert_called_once()
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_cleanup(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
"""Test that re_init_pipeline properly cleans up old pipeline."""
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomPipelineWorkerExtension,
)
# Setup mock model_runner with pipeline
mock_model_runner = Mock()
mock_pipeline = Mock()
mock_model_runner.pipeline = mock_pipeline
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
# Call re_init_pipeline
wrapper.worker.re_init_pipeline(custom_args)
# Verify cleanup was performed
mock_gc_collect.assert_called_once()
mock_empty_cache.assert_called_once()
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_none_pipeline(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
"""Test re_init_pipeline when pipeline is None."""
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomPipelineWorkerExtension,
)
# Setup mock model_runner with None pipeline
mock_model_runner = Mock()
mock_model_runner.pipeline = None
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
# Should not raise an error
wrapper.worker.re_init_pipeline(custom_args)
# Verify load_model was still called
wrapper.worker.load_model.assert_called_once()
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_custom_pipeline_args_initialization(self, mock_worker_init, mock_od_config):
"""Test initialization with custom_pipeline_args calls re_init_pipeline."""
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
with patch.object(DiffusionWorker, "__init__", return_value=None):
with patch.object(WorkerWrapperBase, "_prepare_worker_class") as mock_prepare:
# Create a mock worker class with re_init_pipeline
mock_worker_class = Mock()
mock_worker_instance = Mock()
mock_worker_instance.re_init_pipeline = Mock()
mock_worker_class.return_value = mock_worker_instance
mock_prepare.return_value = mock_worker_class
_ = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
custom_pipeline_args=custom_args,
)
# Verify re_init_pipeline was called with custom_pipeline_args
mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_custom_pipeline_with_explicit_extension(self, mock_worker_init, mock_od_config):
"""Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""
class CustomExtension:
def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]):
return "custom_re_init_pipeline"
def custom_extension_method(self):
return "custom_extension_method"
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomExtension,
custom_pipeline_args=custom_args,
)
# Should still have the explicitly provided extension
assert CustomExtension in wrapper.worker.__class__.__bases__
assert hasattr(wrapper.worker, "custom_extension_method")
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_multiple_calls(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
"""Test calling re_init_pipeline multiple times."""
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomPipelineWorkerExtension,
)
# Setup mock model_runner
mock_model_runner = Mock()
mock_pipeline1 = Mock()
mock_pipeline2 = Mock()
mock_model_runner.pipeline = mock_pipeline1
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
# First call
custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
wrapper.worker.re_init_pipeline(custom_args1)
# Update pipeline for second call
mock_model_runner.pipeline = mock_pipeline2
# Second call
custom_args2 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
wrapper.worker.re_init_pipeline(custom_args2)
# Verify load_model was called twice with different pipelines
assert wrapper.worker.load_model.call_count == 2
assert wrapper.worker.init_lora_manager.call_count == 2
vllm_omni/diffusion/data.py
+9
View File
@@ -332,6 +332,15 @@ class OmniDiffusionConfig:
# TODO: do not hard code
master_port: int | None = None
# Worker extension class for custom functionality
worker_extension_cls: str | None = None
# Custom pipeline arguments for custom pipelines
custom_pipeline_args: dict[str, Any] | None = None
# Diffusion model loading format
diffusion_load_format: str = "default" # "default", "custom_pipeline", "dummy"
# http server endpoint config, would be ignored in local mode
host: str | None = None
port: int | None = None
vllm_omni/diffusion/executor/multiproc_executor.py
+6
View File
@@ -77,6 +77,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
mp.set_start_method("spawn", force=True)
processes = []
# Extract worker_extension_cls and custom_pipeline_args from od_config
worker_extension_cls = od_config.worker_extension_cls
custom_pipeline_args = getattr(od_config, "custom_pipeline_args", None)
# Launch all worker processes
scheduler_pipe_readers = []
scheduler_pipe_writers = []
@@ -91,6 +95,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
od_config,
writer,
broadcast_handle,
worker_extension_cls,
custom_pipeline_args,
),
name=f"DiffusionWorker-{i}",
daemon=True,
vllm_omni/diffusion/model_loader/diffusers_loader.py
+13 -3
View File
@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import (
maybe_download_from_modelscope,
safetensors_weights_iterator,
)
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.utils.torch_utils import set_default_torch_dtype
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -208,13 +209,22 @@ class DiffusersPipelineLoader:
allow_patterns_overrides=None,
)
def load_model(self, od_config: OmniDiffusionConfig, load_device: str) -> nn.Module:
def load_model(
self,
od_config: OmniDiffusionConfig,
load_device: str,
load_format: str = "default",
custom_pipeline_name: str | None = None,
) -> nn.Module:
"""Load a model with the given configurations."""
target_device = torch.device(load_device)
with set_default_torch_dtype(od_config.dtype):
with target_device:
model = initialize_model(od_config)
if load_format == "default":
model = initialize_model(od_config)
elif load_format == "custom_pipeline":
model_cls = resolve_obj_by_qualname(custom_pipeline_name)
model = model_cls(od_config=od_config)
logger.debug("Loading weights on %s ...", load_device)
# Quantization does not happen in `load_weights` but after it
self.load_weights(model)
vllm_omni/diffusion/worker/diffusion_model_runner.py
+14
View File
@@ -70,6 +70,8 @@ class DiffusionModelRunner:
def load_model(
self,
memory_pool_context_fn: callable | None = None,
load_format: str | None = None,
custom_pipeline_name: str | None = None,
) -> None:
"""
Load the diffusion model, apply compilation and offloading.
@@ -77,7 +79,17 @@ class DiffusionModelRunner:
Args:
memory_pool_context_fn: Optional function that returns a context manager
for memory pool allocation (used for sleep mode).
load_format: Format for loading model weights. Supported formats:
- "default" (default): Automatically detect and use the default format based on configuration
- "custom_pipeline": Init model from a custom pipeline class specified by `custom_pipeline_name`
- "dummy": Skip actual weight loading, useful for testing and custom pipelines that
don't require default weights.
custom_pipeline_name: Optional custom pipeline class name to use.
"""
if load_format == "dummy":
return
load_device = (
"cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device)
)
@@ -97,6 +109,8 @@ class DiffusionModelRunner:
self.pipeline = model_loader.load_model(
od_config=self.od_config,
load_device=load_device,
load_format=load_format,
custom_pipeline_name=custom_pipeline_name,
)
time_after_load = time.perf_counter()
vllm_omni/diffusion/worker/diffusion_worker.py
+275 -19
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Diffusion Worker for vLLM-Omni.
@@ -7,15 +8,19 @@ Handles GPU infrastructure initialization and delegates model operations
to DiffusionModelRunner.
"""
import gc
import multiprocessing as mp
import os
from collections.abc import Iterable
from contextlib import AbstractContextManager, nullcontext
from typing import Any
import torch
import zmq
from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm.logger import init_logger
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.utils.mem_utils import GiB_bytes
from vllm.v1.worker.workspace import init_workspace_manager
@@ -67,6 +72,15 @@ class DiffusionWorker:
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
self.lora_manager: DiffusionLoRAManager | None = None
self.init_device()
# Create model runner
self.model_runner = DiffusionModelRunner(
vllm_config=self.vllm_config,
od_config=self.od_config,
device=self.device,
)
self.load_model(load_format=self.od_config.diffusion_load_format)
self.init_lora_manager()
logger.info(f"Worker {self.rank}: Initialization complete.")
def init_device(self) -> None:
"""Initialize the device and distributed environment."""
@@ -110,16 +124,22 @@ class DiffusionWorker:
)
init_workspace_manager(self.device)
# Create model runner and load model
self.model_runner = DiffusionModelRunner(
vllm_config=self.vllm_config,
od_config=self.od_config,
device=self.device,
)
def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None:
"""Load the diffusion model using DiffusionModelRunner."""
with (
set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
set_current_vllm_config(self.vllm_config),
):
self.model_runner.load_model(
memory_pool_context_fn=self._maybe_get_memory_pool_context,
load_format=load_format,
custom_pipeline_name=custom_pipeline_name,
)
assert self.model_runner.pipeline is not None
def init_lora_manager(self) -> None:
"""Initialize the LoRA manager for this worker."""
if self.model_runner.pipeline is None:
return
self.lora_manager = DiffusionLoRAManager(
pipeline=self.model_runner.pipeline,
device=self.device,
@@ -128,7 +148,6 @@ class DiffusionWorker:
lora_path=self.od_config.lora_path,
lora_scale=self.od_config.lora_scale,
)
logger.info(f"Worker {self.rank}: Initialization complete.")
def generate(self, request: OmniDiffusionRequest) -> DiffusionOutput:
"""Generate output for the given requests."""
@@ -247,6 +266,32 @@ class DiffusionWorker:
destroy_distributed_env()
class CustomPipelineWorkerExtension:
def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]) -> None:
"""
Re-initialize the pipeline with custom arguments.
Args:
custom_pipeline_args: Dictionary of arguments for custom pipeline initialization
"""
# Clean up old pipeline
if self.model_runner.pipeline is not None:
del self.model_runner.pipeline
gc.collect()
torch.cuda.empty_cache()
# Get custom pipeline class name
custom_pipeline_name = custom_pipeline_args["pipeline_class"]
# Use the DiffusionWorker's load_model method which handles the forward context
self.load_model(
load_format="custom_pipeline",
custom_pipeline_name=custom_pipeline_name,
)
self.init_lora_manager()
class WorkerProc:
"""Wrapper that runs one Worker in a separate process."""
@@ -255,8 +300,11 @@ class WorkerProc:
od_config: OmniDiffusionConfig,
gpu_id: int,
broadcast_handle,
worker_extension_cls: str | None = None,
custom_pipeline_args: dict[str, Any] | None = None,
):
self.od_config = od_config
self.gpu_id = gpu_id
# Inter-process Communication
self.context = zmq.Context(io_threads=2)
@@ -274,17 +322,26 @@ class WorkerProc:
logger.info(f"Worker {gpu_id} created result MessageQueue")
assert od_config.master_port is not None
self.worker = self._create_worker(gpu_id, od_config)
self.gpu_id = gpu_id
# Create worker using WorkerWrapperBase for extension support
self.worker = self._create_worker(gpu_id, od_config, worker_extension_cls, custom_pipeline_args)
self._running = True
def _create_worker(self, gpu_id: int, od_config: OmniDiffusionConfig) -> DiffusionWorker:
def _create_worker(
self,
gpu_id: int,
od_config: OmniDiffusionConfig,
worker_extension_cls: str | None,
custom_pipeline_args: dict[str, Any] | None = None,
) -> DiffusionWorker:
"""Create a worker instance. Override in subclasses for different worker types."""
return DiffusionWorker(
local_rank=gpu_id,
rank=gpu_id,
wrapper = WorkerWrapperBase(
gpu_id=gpu_id,
od_config=od_config,
worker_extension_cls=worker_extension_cls,
custom_pipeline_args=custom_pipeline_args,
)
return wrapper
def return_result(self, output: DiffusionOutput):
"""Reply to client, only on rank 0."""
@@ -310,11 +367,8 @@ class WorkerProc:
return None, False
try:
if isinstance(method, str):
func = getattr(self.worker, method)
result = func(*args, **kwargs)
else:
result = method(self.worker, *args, **kwargs)
# Use execute_method from WorkerWrapperBase for consistent method resolution
result = self.worker.execute_method(method, *args, **kwargs)
return result, should_reply
except Exception as e:
logger.error(f"Error executing RPC: {e}", exc_info=True)
@@ -385,6 +439,8 @@ class WorkerProc:
od_config: OmniDiffusionConfig,
pipe_writer: mp.connection.Connection,
broadcast_handle,
worker_extension_cls: str | None = None,
custom_pipeline_args: dict[str, Any] | None = None,
) -> None:
"""Worker initialization and execution loops."""
from vllm_omni.plugins import load_omni_general_plugins
@@ -394,6 +450,8 @@ class WorkerProc:
od_config,
gpu_id=rank,
broadcast_handle=broadcast_handle,
worker_extension_cls=worker_extension_cls,
custom_pipeline_args=custom_pipeline_args,
)
logger.info(f"Worker {rank}: Scheduler loop started.")
pipe_writer.send(
@@ -404,3 +462,201 @@ class WorkerProc:
)
worker_proc.worker_busy_loop()
logger.info(f"Worker {rank}: Shutdown complete.")
class WorkerWrapperBase:
"""
Wrapper base class that creates DiffusionWorker with optional worker_extension_cls support.
This enables dynamic inheritance for DiffusionWorker to extend with custom functionality.
"""
def __init__(
self,
gpu_id: int,
od_config: OmniDiffusionConfig,
base_worker_class: type = DiffusionWorker,
worker_extension_cls: str | None = None,
custom_pipeline_args: dict[str, Any] | None = None,
):
"""
Initialize WorkerWrapperBase with support for worker extensions.
Args:
gpu_id: GPU device ID
od_config: OmniDiffusionConfig configuration
worker_extension_cls: Optional qualified name of worker extension class
custom_pipeline_args: Optional arguments for custom pipeline initialization
"""
self.gpu_id = gpu_id
self.od_config = od_config
self.base_worker_class = base_worker_class
self.worker_extension_cls = worker_extension_cls
self.custom_pipeline_args = custom_pipeline_args
# Prepare worker class with extension support
worker_class = self._prepare_worker_class()
# Create the actual worker instance
self.worker = worker_class(
local_rank=gpu_id,
rank=gpu_id,
od_config=od_config,
)
# Re-initialize pipeline with custom pipeline if provided
if self.custom_pipeline_args is not None:
self.worker.re_init_pipeline(self.custom_pipeline_args)
def _prepare_worker_class(self) -> type:
"""
Prepare the worker class with optional extension.
Dynamically extends GPUWorker with worker_extension_cls if provided.
Returns:
The worker class (potentially extended)
"""
worker_class = self.base_worker_class
# If custom_pipeline_args is provided, use CustomPipelineWorkerExtension
if self.custom_pipeline_args is not None:
# Set worker_extension_cls to CustomPipelineWorkerExtension if not already set
if self.worker_extension_cls is None:
self.worker_extension_cls = CustomPipelineWorkerExtension
if self.worker_extension_cls:
if isinstance(self.worker_extension_cls, str):
worker_extension_cls = resolve_obj_by_qualname(self.worker_extension_cls)
else:
worker_extension_cls = self.worker_extension_cls
extended_calls = []
if worker_extension_cls not in worker_class.__bases__:
# Check for conflicts between worker and extension
for attr in dir(worker_extension_cls):
if attr.startswith("__"):
continue
if hasattr(worker_class, attr):
logger.warning(
f"Worker class {worker_class} already has attribute "
f"{attr}, which may conflict with worker extension "
f"class {worker_extension_cls}."
)
if callable(getattr(worker_extension_cls, attr)):
extended_calls.append(attr)
# Dynamically inherit the worker extension class
class_name = f"{worker_class.__name__}With{worker_extension_cls.__name__}"
worker_class = type(class_name, (worker_extension_cls, worker_class), {})
logger.info(
"Created extended worker class %s from %s for extended calls %s",
class_name,
worker_extension_cls,
extended_calls,
)
return worker_class
def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
"""
Generate output for the given requests.
Args:
requests: List of diffusion requests
Returns:
DiffusionOutput with generated results
"""
return self.worker.generate(requests)
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
"""
Execute a forward pass.
Args:
reqs: List of diffusion requests
od_config: OmniDiffusionConfig configuration
Returns:
DiffusionOutput with generated results
"""
return self.worker.execute_model(reqs, od_config)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""
Load model weights.
Args:
weights: Iterable of (name, tensor) tuples
Returns:
Set of loaded weight names
"""
return self.worker.load_weights(weights)
def sleep(self, level: int = 1) -> bool:
"""
Put the worker to sleep. The worker should not process any requests.
The caller should guarantee that no requests are being processed
during the sleep period, before `wake_up` is called.
Args:
level: The sleep level. Level 1 sleep will offload the model
weights and discard the kv cache.
Currently only support level 1.
Returns:
True on success
"""
return self.worker.sleep(level)
def wake_up(self, tags: list[str] | None = None) -> bool:
"""
Wake up the worker from sleep mode. See the sleep function
method for more details.
Args:
tags: An optional list of tags to reallocate the worker memory
for specific memory allocations. Values must be in
`("weights")`. If None, all memory is reallocated.
wake_up should be called with all tags (or None) before the
worker is used again.
Returns:
True on success
"""
return self.worker.wake_up(tags)
def shutdown(self) -> None:
"""Shutdown the worker and cleanup resources."""
return self.worker.shutdown()
def execute_method(self, method: str | bytes, *args, **kwargs) -> Any:
"""
Execute a method on the worker.
Args:
method: Method name (str) or serialized callable (bytes)
Returns:
Result of the method execution (type depends on the method)
Raises:
Exception: If method execution fails
"""
try:
# Method resolution order:
# 1. If method is defined in this class, it will be called directly
# 2. Otherwise, since we define `__getattr__` and redirect attribute
# query to `self.worker`, the method will be called on the worker
assert isinstance(method, str), "Method must be str"
func = getattr(self.worker, method)
return func(*args, **kwargs)
except Exception as e:
msg = f"Error executing method {method!r}. This might cause issues in distributed execution."
logger.exception(msg)
raise e
def __getattr__(self, attr: str):
"""Delegate attribute access to the wrapped worker."""
return getattr(self.worker, attr)
vllm_omni/entrypoints/async_omni.py
+2
View File
@@ -163,6 +163,8 @@ class AsyncOmni(OmniBase):
"enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
"enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
"enforce_eager": kwargs.get("enforce_eager", False),
"diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
"custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
},
"final_output": True,
"final_output_type": "image",