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. diffusion_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

diffusion_load_format parameter¶

The diffusion_load_format parameter controls the initial diffusion model load. The relevant values are:

• "default": Standard model loading using the model registry (default behavior)
• "dummy": Skip the initial model load; use this with custom_pipeline_args["pipeline_class"] for the custom pipeline workflow below
• "diffusers": Load through the Hugging Face Diffusers adapter

Location: vllm_omni/diffusion/data.py (OmniDiffusionConfig)

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:

# custom_pipeline.py
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
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: DiffusionRequestBatch):
        # Optionally customize sampling parameters
        actual_num_steps = req.sampling_params.num_inference_steps or 50
        req.sampling_params.num_inference_steps = actual_num_steps

        # Call parent's forward to get normal output
        output = super().forward(req=req)

        # Add custom trajectory data
        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:

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]
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:

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:

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
)