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

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

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
+)
+```