From 99945947827efd1f58e589d8a791b1d7e873a3bc Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Thu, 12 Feb 2026 19:23:51 +0800 Subject: [PATCH] [Misc] Support WorkerWrapperBase and CustomPipeline for Diffusion Worker (#764) Signed-off-by: knlnguyen1802 --- docs/features/custom_pipeline.md | 151 +++++ .../image_to_image/custom_pipeline.py | 80 +++ .../image_to_image/image_edit.py | 262 +++++++++ .../custom_pipeline/image_to_image/run.sh | 6 + tests/diffusion/test_diffusion_worker.py | 2 +- tests/diffusion/test_worker_wrapper_base.py | 540 ++++++++++++++++++ vllm_omni/diffusion/data.py | 9 + .../diffusion/executor/multiproc_executor.py | 6 + .../model_loader/diffusers_loader.py | 16 +- .../worker/diffusion_model_runner.py | 14 + .../diffusion/worker/diffusion_worker.py | 294 +++++++++- vllm_omni/entrypoints/async_omni.py | 2 + 12 files changed, 1359 insertions(+), 23 deletions(-) create mode 100644 docs/features/custom_pipeline.md create mode 100644 examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py create mode 100644 examples/offline_inference/custom_pipeline/image_to_image/image_edit.py create mode 100644 examples/offline_inference/custom_pipeline/image_to_image/run.sh create mode 100644 tests/diffusion/test_worker_wrapper_base.py diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md new file mode 100644 index 000000000..be46b5b14 --- /dev/null +++ b/docs/features/custom_pipeline.md @@ -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 +) +``` diff --git a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py new file mode 100644 index 000000000..b1c74e2b3 --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py @@ -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 diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py new file mode 100644 index 000000000..8ad5cbf9b --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -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()) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/run.sh b/examples/offline_inference/custom_pipeline/image_to_image/run.sh new file mode 100644 index 000000000..1d2706e25 --- /dev/null +++ b/examples/offline_inference/custom_pipeline/image_to_image/run.sh @@ -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 diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 6e57355a2..1c767b98f 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -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() diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py new file mode 100644 index 000000000..2023a89e1 --- /dev/null +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a18411dba..cbfa1b1bc 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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 diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index a421a01ed..f26d89d20 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -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, diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index b61f70b69..7806789b1 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -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) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 8362eef18..134a2ae39 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -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() diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 6f5bd00e6..90fff548e 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -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) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 0be77bf73..369713d7b 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -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",