1088 lines
38 KiB
Markdown
1088 lines
38 KiB
Markdown
# Adding a Diffusion Model
|
||
|
||
This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns.
|
||
|
||
|
||
---
|
||
|
||
## Table of Contents
|
||
|
||
1. [Overview](#overview)
|
||
2. [Directory Structure](#directory-structure)
|
||
3. [Basic Implementation](#basic-implementation)
|
||
4. [Advanced Features](#advanced-features)
|
||
5. [Troubleshooting](#troubleshooting)
|
||
6. [Pull Request Checklist](#pull-request-checklist)
|
||
7. [Reference Implementations](#reference-implementations)
|
||
8. [Summary](#summary)
|
||
|
||
---
|
||
|
||
## Overview
|
||
|
||
vLLM-Omni's diffusion inference follows this architecture:
|
||
|
||
<p align="center">
|
||
<picture>
|
||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-diffusion-flow.png">
|
||
<img alt="Diffusion Flow" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-diffusion-flow.png" width=55%>
|
||
</picture>
|
||
</p>
|
||
|
||
**Key Components:**
|
||
|
||
1. **Request Handling:** User prompts → `OmniDiffusionRequest`
|
||
2. **Diffusion Engine:** Request → Preprocessing (Optional) → Pipeline execution -> Post-processing
|
||
3. **Pipeline Execution:** Request → Encode prompt → Diffusion steps → Vae decode
|
||
|
||
|
||
## Directory Structure
|
||
|
||
Organize your model files following this structure:
|
||
|
||
```
|
||
vllm_omni/
|
||
└── diffusion/
|
||
├── registry.py # ← Register your model here
|
||
├── request.py # Request data structures
|
||
└── models/
|
||
└── your_model_name/ # ← Create this directory
|
||
├── __init__.py # Export pipeline and transformer
|
||
├── pipeline_xxx.py # Pipeline implementation
|
||
└── xxx_transformer.py # Transformer implementation
|
||
```
|
||
|
||
**Naming Conventions:**
|
||
|
||
- **Model directory:** `your_model_name` (lowercase, underscores), e.g., `qwen_image`, `flux`, `longcat_image`, `wan2_2`
|
||
- **Pipeline file:** `pipeline_xxx.py` where `xxx` describes the task, e.g., `pipeline_qwen_image.py`, `pipeline_qwen_image_edit.py`
|
||
- **Transformer file:** `xxx_transformer.py` matching transformer class name, e.g., `qwen_image_transformer.py`, `flux_transformer.py`
|
||
|
||
---
|
||
|
||
## Basic Implementation
|
||
|
||
This section covers the minimal steps to get a model working in vLLM-Omni with basic features (online/offline serving, batch requests).
|
||
|
||
### Step 1: Adapt Transformer Model
|
||
|
||
The transformer is the core denoising network. Start by copying the transformer implementation from Diffusers and making these adaptations.
|
||
|
||
|
||
#### 1.1: Remove Diffusers Mixins
|
||
|
||
Diffusers' `Mixin` classes are not needed in vLLM-Omni. Remove them:
|
||
|
||
```diff
|
||
# Before (Diffusers)
|
||
- from diffusers.models.modeling_utils import ModelMixin
|
||
- from diffusers.models.attention_processor import AttentionModuleMixin
|
||
|
||
- class YourModelTransformer2DModel(ModelMixin, AttentionModuleMixin):
|
||
+ class YourModelTransformer2DModel(nn.Module):
|
||
"""Your transformer model."""
|
||
```
|
||
|
||
**Example mixins to remove:**
|
||
|
||
- `ModelMixin` - Weight loading utilities (vLLM-Omni has its own weight loader)
|
||
- `AttentionModuleMixin` - Attention processors (using vLLM-Omni's Attention layer instead)
|
||
- `ConfigMixin` - Config management (not needed)
|
||
- `PeftAdapterMixin` - Parameter efficient finetune utilities (not needed)
|
||
|
||
#### 1.2: Replace Attention Implementation
|
||
|
||
**The most important adaptation:** Replace Diffusers' attention with vLLM-Omni's optimized `Attention` layer.
|
||
|
||
**Before (Diffusers):**
|
||
```python
|
||
from diffusers.models.attention_processor import dispatch_attention_fn
|
||
|
||
class YourAttentionBlock(nn.Module):
|
||
def forward(self, hidden_states, encoder_hidden_states=None, ...):
|
||
...
|
||
hidden_states = dispatch_attention_fn(
|
||
query, key, value,
|
||
attn_mask=attention_mask,
|
||
dropout_p=0.0,
|
||
is_causal=False,
|
||
backend=self._attention_backend,
|
||
)
|
||
```
|
||
|
||
**After (vLLM-Omni):**
|
||
```python
|
||
from vllm_omni.diffusion.attention.layer import Attention
|
||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||
|
||
class YourAttentionBlock(nn.Module):
|
||
def __init__(self, ...):
|
||
super().__init__()
|
||
|
||
# Initialize vLLM-Omni's Attention layer
|
||
self.attn = Attention(
|
||
num_heads=self.num_heads,
|
||
head_size=self.head_dim,
|
||
softmax_scale=1.0 / (self.head_dim ** 0.5),
|
||
causal=False, # Diffusion models typically use bidirectional attention
|
||
num_kv_heads=self.num_kv_heads,
|
||
)
|
||
|
||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...):
|
||
...
|
||
# Create attention metadata
|
||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||
hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
|
||
|
||
```
|
||
|
||
**Key Points:**
|
||
|
||
- **Attention layer initialization:** Done in `__init__`, not per-forward
|
||
- **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape
|
||
- **AttentionMetadata:** Wraps attention mask and other metadata
|
||
|
||
**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models.
|
||
|
||
#### 1.3: Replace Imports and Utilities
|
||
|
||
**Logger:**
|
||
```diff
|
||
- from diffusers.utils import logging
|
||
- logger = logging.get_logger(__name__)
|
||
|
||
+ from vllm.logger import init_logger
|
||
+ logger = init_logger(__name__)
|
||
```
|
||
|
||
**Custom layers from vLLM and vLLM-Omni (if needed):**
|
||
|
||
```python
|
||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
|
||
```
|
||
|
||
#### 1.4: Remove Training-Only Code
|
||
|
||
Remove code that's only needed for training:
|
||
|
||
```diff
|
||
# Remove gradient checkpointing
|
||
- if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||
- hidden_states = torch.utils.checkpoint.checkpoint(
|
||
- self._forward_block, hidden_states, ...
|
||
- )
|
||
- else:
|
||
- hidden_states = self._forward_block(hidden_states, ...)
|
||
+ hidden_states = self._forward_block(hidden_states, ...)
|
||
|
||
# Remove training-specific attributes
|
||
- self.gradient_checkpointing = False
|
||
|
||
# Remove dropout (set to 0 or remove)
|
||
- self.dropout = nn.Dropout(dropout_prob)
|
||
+ # Removed dropout for inference
|
||
```
|
||
|
||
#### 1.5: Add Configuration Support
|
||
|
||
Add support for vLLM-Omni's `OmniDiffusionConfig`:
|
||
|
||
```python
|
||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||
|
||
class YourModelTransformer2DModel(nn.Module):
|
||
def __init__(
|
||
self,
|
||
*,
|
||
od_config: OmniDiffusionConfig | None = None, # ← Add vLLM-Omni config
|
||
# ... other model-specific parameters
|
||
num_layers: int = 28,
|
||
hidden_size: int = 3072,
|
||
num_heads: int = 24,
|
||
**kwargs,
|
||
):
|
||
super().__init__()
|
||
|
||
# Store config
|
||
self.od_config = od_config
|
||
self.parallel_config = od_config.parallel_config if od_config else None
|
||
|
||
# Model architecture
|
||
self.num_layers = num_layers
|
||
self.hidden_size = hidden_size
|
||
# ... initialize layers
|
||
```
|
||
|
||
### Step 2: Adapt Pipeline
|
||
|
||
The pipeline orchestrates the full generation process (text encoding, denoising loop, VAE decoding). Adapt it from Diffusers format to vLLM-Omni's interface.
|
||
|
||
#### 2.1: Remove Diffusers Inheritance
|
||
|
||
**Remove Diffusers base classes:**
|
||
```diff
|
||
- from diffusers import DiffusionPipeline
|
||
- from diffusers.loaders import LoraLoaderMixin
|
||
|
||
- class YourModelPipeline(DiffusionPipeline, LoraLoaderMixin):
|
||
+ class YourModelPipeline(nn.Module):
|
||
"""Your model pipeline for vLLM-Omni."""
|
||
```
|
||
|
||
#### 2.2: Adapt `__init__` Method
|
||
|
||
**Before (Diffusers):**
|
||
```python
|
||
class YourModelPipeline(DiffusionPipeline):
|
||
def __init__(
|
||
self,
|
||
vae: AutoencoderKL,
|
||
text_encoder: CLIPTextModel,
|
||
tokenizer: CLIPTokenizer,
|
||
transformer: YourTransformer,
|
||
scheduler: FlowMatchScheduler,
|
||
):
|
||
super().__init__()
|
||
self.register_modules(
|
||
vae=vae,
|
||
text_encoder=text_encoder,
|
||
tokenizer=tokenizer,
|
||
transformer=transformer,
|
||
scheduler=scheduler,
|
||
)
|
||
```
|
||
|
||
**After (vLLM-Omni):**
|
||
```python
|
||
import os
|
||
from diffusers import AutoencoderKL
|
||
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
||
from transformers import CLIPTextModel, CLIPTokenizer
|
||
|
||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
|
||
from vllm_omni.diffusion.models.your_model_name.your_model_transformer import (
|
||
YourModelTransformer2DModel,
|
||
)
|
||
|
||
|
||
class YourModelPipeline(nn.Module):
|
||
def __init__(
|
||
self,
|
||
*,
|
||
od_config: OmniDiffusionConfig,
|
||
prefix: str = "",
|
||
):
|
||
super().__init__()
|
||
self.od_config = od_config
|
||
self.parallel_config = od_config.parallel_config
|
||
self.device = get_local_device()
|
||
model = od_config.model
|
||
local_files_only = os.path.exists(model)
|
||
|
||
# Load components from checkpoint
|
||
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||
model, subfolder="scheduler", local_files_only=local_files_only)
|
||
self.text_encoder = CLIPTextModel.from_pretrained(
|
||
model, subfolder="text_encoder", local_files_only=local_files_only).to(self.device)
|
||
self.tokenizer = CLIPTokenizer.from_pretrained(
|
||
model, subfolder="tokenizer", local_files_only=local_files_only)
|
||
self.vae = AutoencoderKL.from_pretrained(
|
||
model, subfolder="vae", local_files_only=local_files_only).to(self.device)
|
||
|
||
# Initialize transformer with vLLM-Omni config
|
||
transformer_kwargs = get_transformer_config_kwargs(
|
||
od_config.tf_model_config, YourModelTransformer2DModel)
|
||
self.transformer = YourModelTransformer2DModel(
|
||
od_config=od_config, **transformer_kwargs)
|
||
|
||
# Store VAE scale factor for latent space conversions
|
||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||
self.default_sample_size = 128 # Default latent size
|
||
```
|
||
|
||
**Key Changes:**
|
||
|
||
1. **`od_config` parameter:** All configuration through `OmniDiffusionConfig`
|
||
2. **Manual component loading:** No `register_modules()`, load each component explicitly
|
||
3. **Local files support:** Check `os.path.exists(model)` for local checkpoints
|
||
4. **Transformer with config:** Pass `od_config` to transformer constructor
|
||
|
||
#### 2.3: Adapt `__call__` → `forward` Method
|
||
|
||
**Change signature:**
|
||
|
||
```diff
|
||
- @torch.no_grad()
|
||
- def __call__(
|
||
+ def forward(
|
||
self,
|
||
+ req: OmniDiffusionRequest, # ← Add request parameter here
|
||
- ):
|
||
+ ) -> DiffusionOutput: # ← Add return type
|
||
```
|
||
|
||
[`OmniDiffusionRequest`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/request/#vllm_omni.diffusion.request.OmniDiffusionRequest) is a dataclass that contains the **prompts** and **sampling parameters** [`OmniDiffusionSamplingParams`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/inputs/data/#vllm_omni.inputs.data.OmniDiffusionSamplingParams) for the diffusion pipeline execution. It also contains a request_id for other components to trace this request and its outputs.
|
||
|
||
See some parameters in `OmniDiffusionSamplingParams` as follows:
|
||
|
||
| parameters | type |value | function |
|
||
|:---:|:---:|:---:|:---:|
|
||
| `num_inference_steps` | `int` | 50 | The number of diffusion steps during inference|
|
||
| `guidance_scale` | `float` | 0.0 | The classifier free guidance scale |
|
||
| `width` and `height` | `int` | None | The width and height of the generated image |
|
||
|
||
**Extract parameters from request:**
|
||
|
||
```python
|
||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||
from vllm_omni.diffusion.data import DiffusionOutput
|
||
|
||
def forward(
|
||
self,
|
||
req: OmniDiffusionRequest,
|
||
) -> DiffusionOutput:
|
||
# Extract prompts from request
|
||
if req.prompts is not None:
|
||
prompt = [
|
||
p if isinstance(p, str) else (p.get("prompt") or "")
|
||
for p in req.prompts
|
||
]
|
||
|
||
# Extract sampling parameters
|
||
sampling_params = req.sampling_params
|
||
num_inference_steps = sampling_params.num_inference_steps or 50
|
||
guidance_scale = sampling_params.guidance_scale or 7.5
|
||
height = sampling_params.height or (self.default_sample_size * self.vae_scale_factor)
|
||
width = sampling_params.width or (self.default_sample_size * self.vae_scale_factor)
|
||
|
||
# For image editing pipelines, extract images from multi_modal_data
|
||
if hasattr(req, 'multi_modal_data') and req.multi_modal_data:
|
||
input_images = req.multi_modal_data.get('image', [])
|
||
|
||
# ... rest of generation logic
|
||
```
|
||
|
||
For an image editing model, an example `OmniDiffusionRequest` is like:
|
||
```python
|
||
{
|
||
"prompt": "turn this cat to a dog",
|
||
"multi_modal_data": {"image": input_image}
|
||
},
|
||
```
|
||
|
||
**Wrap output:**
|
||
|
||
```diff
|
||
# Generate images
|
||
images = self.vae.decode(latents)[0]
|
||
|
||
- return {"images": images}
|
||
+ return DiffusionOutput(output=images)
|
||
```
|
||
|
||
#### 2.4: Extract Pre/Post-Processing Functions
|
||
|
||
vLLM-Omni separates image processing from the main pipeline for better modularity.
|
||
|
||
**Post-processing function (required):**
|
||
```python
|
||
def get_your_model_post_process_func(
|
||
od_config: OmniDiffusionConfig,
|
||
):
|
||
"""
|
||
Create post-processing function for your model.
|
||
|
||
Returns a function that converts latents to images.
|
||
"""
|
||
from diffusers.image_processor import VaeImageProcessor
|
||
import json
|
||
|
||
# Load VAE config to get scale factor
|
||
model_path = od_config.model
|
||
if not os.path.exists(model_path):
|
||
from vllm_omni.diffusion.model_loader.utils import download_weights_from_hf_specific
|
||
model_path = download_weights_from_hf_specific(model_path, None, ["*"])
|
||
|
||
vae_config_path = os.path.join(model_path, "vae/config.json")
|
||
with open(vae_config_path) as f:
|
||
vae_config = json.load(f)
|
||
vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1)
|
||
|
||
# Create image processor
|
||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
|
||
|
||
def post_process_func(images: torch.Tensor):
|
||
return image_processor.postprocess(images, output_type="pil")
|
||
|
||
return post_process_func
|
||
```
|
||
|
||
**Pre-processing function (for image editing pipelines):**
|
||
|
||
```python
|
||
def get_your_model_pre_process_func(
|
||
od_config: OmniDiffusionConfig,
|
||
):
|
||
"""
|
||
Create pre-processing function for image editing.
|
||
|
||
Returns a function that prepares input images.
|
||
"""
|
||
from PIL import Image
|
||
from diffusers.image_processor import VaeImageProcessor
|
||
|
||
# Load VAE config
|
||
# ... (similar to post_process_func)
|
||
|
||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
|
||
|
||
def pre_process_func(
|
||
request: OmniDiffusionRequest,
|
||
):
|
||
for i, prompt in enumerate(request.prompts):
|
||
multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
|
||
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
|
||
# image pre-processing
|
||
# after pre-processing, update the request attributes
|
||
...
|
||
return request
|
||
|
||
return pre_process_func
|
||
```
|
||
|
||
#### 2.5: Add Weight Loading Support
|
||
|
||
Add methods for automatic weight downloading and loading:
|
||
|
||
```python
|
||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||
|
||
class YourModelPipeline(nn.Module):
|
||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||
super().__init__()
|
||
# ... initialization code
|
||
|
||
# Define weight sources for automatic loading
|
||
self.weights_sources = [
|
||
DiffusersPipelineLoader.ComponentSource(
|
||
model_or_path=od_config.model,
|
||
subfolder="transformer",
|
||
revision=None,
|
||
prefix="transformer.",
|
||
fall_back_to_pt=True,
|
||
)
|
||
]
|
||
|
||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||
"""
|
||
Customize the weight loading behavior, such as filter weights name.
|
||
|
||
Args:
|
||
weights: Iterable of (param_name, param_tensor) tuples
|
||
|
||
Returns:
|
||
Set of loaded parameter names
|
||
"""
|
||
loader = AutoWeightsLoader(self)
|
||
return loader.load_weights(weights)
|
||
```
|
||
|
||
### Step 3: Register Model
|
||
|
||
Register your model in `vllm_omni/diffusion/registry.py` so vLLM-Omni can discover and load it.
|
||
|
||
#### 3.1: Register Pipeline Class
|
||
|
||
```python
|
||
# vllm_omni/diffusion/registry.py
|
||
|
||
_DIFFUSION_MODELS = {
|
||
# Format: "PipelineClassName": (module_folder, module_file, class_name)
|
||
|
||
# Existing models
|
||
"QwenImagePipeline": ("qwen_image", "pipeline_qwen_image", "QwenImagePipeline"),
|
||
"FluxPipeline": ("flux", "pipeline_flux", "FluxPipeline"),
|
||
|
||
# Add your model
|
||
"YourModelPipeline": (
|
||
"your_model_name", # Module folder name
|
||
"pipeline_your_model", # Python file name (without .py)
|
||
"YourModelPipeline", # Pipeline class name
|
||
),
|
||
}
|
||
```
|
||
|
||
#### 3.2: Register Pre/Post-Processing Function
|
||
|
||
```python
|
||
# vllm_omni/diffusion/registry.py
|
||
_DIFFUSION_PRE_PROCESS_FUNCS = {
|
||
# arch: pre_process_func
|
||
# `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
|
||
# where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
|
||
"GlmImagePipeline": "get_glm_image_pre_process_func",
|
||
"QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
|
||
|
||
# Add your model
|
||
"YourModelPipeline": "get_your_model_pre_process_func", # Optional
|
||
}
|
||
_DIFFUSION_POST_PROCESS_FUNCS = {
|
||
# Format: "PipelineClassName": "function_name"
|
||
|
||
# Existing models
|
||
"QwenImagePipeline": "get_qwen_image_post_process_func",
|
||
"FluxPipeline": "get_flux_post_process_func",
|
||
|
||
# Add your model
|
||
"YourModelPipeline": "get_your_model_post_process_func",
|
||
}
|
||
```
|
||
|
||
|
||
#### 3.3: Export from Module
|
||
|
||
Create/update `__init__.py` to export your classes:
|
||
|
||
```python
|
||
# vllm_omni/diffusion/models/your_model_name/__init__.py
|
||
|
||
from .pipeline_your_model import (
|
||
YourModelPipeline,
|
||
get_your_model_post_process_func,
|
||
)
|
||
from .your_model_transformer import YourModelTransformer2DModel
|
||
|
||
__all__ = [
|
||
"YourModelPipeline",
|
||
"YourModelTransformer2DModel",
|
||
"get_your_model_post_process_func",
|
||
]
|
||
```
|
||
|
||
---
|
||
|
||
### Step 4: Add Example Script
|
||
|
||
|
||
If your model is one of Text-to-Image, Text-to-Audio, Text-to-Video, Image-to-Image, Image-to-Video models, you can simply try one of the following offline inference scripts to run your model:
|
||
|
||
| Model Category | Offline Inference Script |
|
||
|---|---|
|
||
| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
|
||
| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
|
||
| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
|
||
| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
|
||
| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
|
||
|
||
|
||
If new CLI arguments need to be added, please edit the offline inference script corresponding to your model category from the table above, and update the example inference script in its corresponding document file (e.g., `examples/offline_inference/text_to_video/text_to_video.md`).
|
||
|
||
For online inference, all the supported tasks are listed in `docs/user_guide/examples/online_serving/`. If your model falls into these categories, please check the corresponding documentation in this folder and the example at `examples/online_serving/TASK_NAME`. Update them accordingly if needed.
|
||
|
||
---
|
||
|
||
If your model is an Omni (understanding and generation) model, please follow the steps below.
|
||
|
||
#### 4.1: Create Example File
|
||
|
||
Taking **BAGEL** model as examples for both offline and online:
|
||
|
||
- Offline: `examples/offline_inference/bagel/`
|
||
- Online: `examples/online_serving/bagel/`
|
||
|
||
Add **two example folders** for your model:
|
||
|
||
```bash
|
||
mkdir -p examples/offline_inference/your_model_name
|
||
mkdir -p examples/online_serving/your_model_name
|
||
```
|
||
|
||
**Offline (recommended minimum):** create `examples/offline_inference/your_model_name/end2end.py` and a README.
|
||
|
||
- Script: `examples/offline_inference/your_model_name/end2end.py`
|
||
- Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.)
|
||
- Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only)
|
||
- Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`)
|
||
- Doc: `examples/offline_inference/your_model_name/README.md`
|
||
- Include at least one runnable command, e.g.:
|
||
|
||
```bash
|
||
cd examples/offline_inference/your_model_name
|
||
python end2end.py --model your-org/your-model-name --modality text2img --prompts "A cute cat"
|
||
```
|
||
|
||
#### 4.2: Add Online Serving Example (OpenAI-Compatible)
|
||
|
||
Mirror BAGEL’s online serving setup:
|
||
|
||
- Server launcher: `examples/online_serving/your_model_name/run_server.sh`
|
||
- Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed)
|
||
- Client: `examples/online_serving/your_model_name/openai_chat_client.py`
|
||
- Send requests to `POST /v1/chat/completions`
|
||
- Support multimodal inputs (e.g., base64 image) if your model needs it
|
||
- Doc: `examples/online_serving/your_model_name/README.md`
|
||
- Include both “launch server” and “send request”:
|
||
|
||
```bash
|
||
# Terminal 1: launch server
|
||
cd examples/online_serving/your_model_name
|
||
bash run_server.sh
|
||
|
||
# Terminal 2: send request
|
||
python openai_chat_client.py --prompt "A cute cat" --modality text2img
|
||
```
|
||
|
||
|
||
### Step 5: Test Your Implementation
|
||
|
||
Before submitting, thoroughly test your implementation.
|
||
|
||
#### 5.1: Performance/Speed Check
|
||
|
||
Manually compare **latency/throughput** and **output quality** against a Diffusers baseline.
|
||
|
||
For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num_inference_steps**, and **guidance settings**, and run multiple trials to reduce randomness. Record the results (and your hardware / driver / CUDA versions) in your PR description.
|
||
|
||
|
||
#### 5.2 Functionality Check in CI
|
||
|
||
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
|
||
|
||
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||
The following tests are required to add:
|
||
|
||
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
|
||
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
|
||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
|
||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
|
||
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
|
||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
|
||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
|
||
- Validation: test that the multimodal output files of your model have the correct shapes.
|
||
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
|
||
- To maximize code reuse, you may refer to
|
||
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
|
||
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
|
||
- Doc: add a concise dostring for each test function.
|
||
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
|
||
|
||
---
|
||
|
||
## Advanced Features
|
||
|
||
Once basic implementation works, add advanced features for better performance.
|
||
|
||
### torch.compile Support
|
||
|
||
Enable automatic compilation for repeated blocks:
|
||
|
||
```python
|
||
# In your_model_transformer.py
|
||
|
||
class YourModelTransformer2DModel(nn.Module):
|
||
# Specify which blocks can be compiled
|
||
_repeated_blocks = ["YourTransformerBlock"] # List of block class names
|
||
|
||
def __init__(self, ...):
|
||
super().__init__()
|
||
# ... initialization
|
||
```
|
||
|
||
vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compile` is available.
|
||
|
||
### Tensor Parallelism
|
||
|
||
See detailed guide: [How to add Tensor Parallel support](../../design/feature/tensor_parallel.md)
|
||
|
||
**Quick setup:**
|
||
|
||
1. Replace Linear layers by various parallel linear layers (e.g., `ColumnParallelLinear`) in vLLM
|
||
2. Check `tp_size` validity: `hidden_dim`, `num_heads`, and `num_kv_heads` must be divisible by `tp_size`
|
||
|
||
**Usage:** Set `tensor_parallel_size` when initializing:
|
||
```python
|
||
omni = Omni(model="your-model", tensor_parallel_size=2)
|
||
```
|
||
|
||
### CFG Parallelism
|
||
|
||
See detailed guide: [How to add CFG-Parallel support](../../design/feature/cfg_parallel.md)
|
||
|
||
**Quick setup:**
|
||
|
||
1. Implement `diffuse()` method
|
||
2. Inherit `CFGParallelMixin` in your pipeline class
|
||
|
||
**Usage:** Set `cfg_parallel_size` when initializing:
|
||
```python
|
||
omni = Omni(model="your-model", cfg_parallel_size=2)
|
||
```
|
||
|
||
### Sequence Parallelism
|
||
|
||
See detailed guide: [How to add Sequence Parallel support](../../design/feature/sequence_parallel.md)
|
||
|
||
**Quick setup:**
|
||
|
||
1. Add `_sp_plan` class attribute to transformer
|
||
2. Specify where to shard/gather tensors
|
||
|
||
**Usage:** Set `ulysses_degree` and `ring_degree` when initializing:
|
||
```python
|
||
omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
|
||
```
|
||
|
||
### Step Execution
|
||
|
||
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
|
||
|
||
Use this only when your pipeline can be split into stable request-scoped and
|
||
step-scoped phases. The reference implementation is
|
||
`QwenImagePipeline`, which maps its request-level `forward()` into:
|
||
|
||
1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
|
||
2. `denoise_step()` for one transformer/noise prediction.
|
||
3. `step_scheduler()` for one scheduler update and `step_index` advance.
|
||
4. `post_decode()` for the final VAE decode.
|
||
|
||
Do not enable `step_execution=True` until those four methods are implemented
|
||
and validated against the request-level path.
|
||
|
||
### Cache Acceleration
|
||
|
||
#### TeaCache
|
||
|
||
See detailed guide: [How to add TeaCache support](../../design/feature/teacache.md)
|
||
|
||
**Quick setup:**
|
||
|
||
1. Write extractor function
|
||
2. Register in `EXTRACTOR_REGISTRY`
|
||
3. Add polynomial coefficients
|
||
|
||
**Usage:** Set `cache_backend` and `cache_config` when initializing:
|
||
```python
|
||
omni = Omni(model="your-model",
|
||
cache_backend="tea_cache",
|
||
cache_config={"rel_l1_thresh": 0.2}
|
||
)
|
||
|
||
```
|
||
|
||
|
||
#### Cache-DiT
|
||
|
||
See detailed guide: [How to add Cache-DiT support](../../design/feature/cache_dit.md)
|
||
|
||
**Quick setup:**
|
||
|
||
- For standard models: Works automatically
|
||
- For complex architectures: Write custom cache config
|
||
|
||
**Usage:** Set `cache_backend` and `cache_config` when initializing:
|
||
```python
|
||
omni = Omni(model="your-model",
|
||
cache_backend="cache_dit",
|
||
cache_config={
|
||
"Fn_compute_blocks": 1,
|
||
"Bn_compute_blocks": 0,
|
||
"max_warmup_steps": 4,
|
||
}
|
||
)
|
||
```
|
||
|
||
### CPU Offload
|
||
|
||
See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md)
|
||
|
||
vLLM-Omni provides two offloading strategies to reduce GPU memory usage:
|
||
|
||
1. **Model-level offload**: Mutual exclusion between DiT and encoders (only one on GPU at a time)
|
||
2. **Layerwise (Blockwise) offload**: Keeps only a single transformer block on GPU at a time with compute-memory overlap
|
||
|
||
**Usage:** Enable offload when initializing:
|
||
```python
|
||
# Model-level offload
|
||
omni = Omni(model="your-model", enable_cpu_offload=True)
|
||
|
||
# Layerwise offload
|
||
omni = Omni(model="your-model", enable_layerwise_offload=True)
|
||
```
|
||
|
||
**To support layerwise offloading:** Define the blocks attribute name in your transformer:
|
||
|
||
```python
|
||
class WanTransformer3DModel(nn.Module):
|
||
_layerwise_offload_blocks_attr = "blocks" # Attribute name containing transformer blocks
|
||
|
||
def __init__(self):
|
||
self.blocks = nn.ModuleList([...]) # Transformer blocks
|
||
```
|
||
|
||
**Note:** Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations.
|
||
|
||
|
||
---
|
||
|
||
### Diffusion Timing (Performance Profiling)
|
||
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
|
||
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
|
||
|
||
!!! info
|
||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||
|
||
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
|
||
|
||
**Enabling Diffusion Timing**
|
||
|
||
|
||
Enable timing by setting:
|
||
```
|
||
vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
|
||
```
|
||
You can optionally specify which modules to profile:
|
||
```
|
||
class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
|
||
def __init__(self, xxx):
|
||
...
|
||
self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
|
||
```
|
||
If not specified, the default targets are used:
|
||
```
|
||
["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
|
||
```
|
||
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
|
||
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
|
||
```python
|
||
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||
|
||
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||
# Optional: Specify custom timing targets
|
||
_PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
od_config: OmniDiffusionConfig,
|
||
prefix: str = "",
|
||
):
|
||
super().__init__()
|
||
self.od_config = od_config
|
||
self.parallel_config = od_config.parallel_config
|
||
# initialize pipeline components
|
||
...
|
||
|
||
# initialize timing profiler
|
||
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
|
||
```
|
||
The mixin dynamically wraps selected methods and records their execution time during inference.
|
||
|
||
If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
|
||
|
||
```diff
|
||
- return DiffusionOutput(output=img)
|
||
+ return DiffusionOutput(
|
||
output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
||
)
|
||
```
|
||
|
||
**Pipeline Design for Timing**
|
||
The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
|
||
|
||
When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
|
||
|
||
For example:
|
||
```
|
||
def forward(self, req: OmniDiffusionRequest):
|
||
prompt_embeds = self.encode_prompt(req)
|
||
latents = self.diffuse(prompt_embeds, req)
|
||
images = self.vae.decode(latents)
|
||
return DiffusionOutput(output=images)
|
||
```
|
||
This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
|
||
|
||
|
||
**Default Profiled Modules**
|
||
|
||
By default, the following pipeline modules are timed:
|
||
```
|
||
vae.encode
|
||
vae.decode
|
||
diffuse
|
||
text_encoder.forward
|
||
tokenizer.forward
|
||
```
|
||
|
||
**Example Output**
|
||
|
||
When enabled, timing logs appear like this:
|
||
```
|
||
[DiffusionTiming] text_encoder.forward took 0.018s
|
||
[DiffusionTiming] diffuse took 2.412s
|
||
[DiffusionTiming] vae.decode took 0.063s
|
||
```
|
||
These measurements help identify bottlenecks during model adaptation and optimization
|
||
|
||
|
||
|
||
## Troubleshooting
|
||
|
||
|
||
**Issue: ImportError when loading model**
|
||
|
||
**Symptoms:** `ModuleNotFoundError` or `ImportError` when calling `Omni(model="your-model")`
|
||
|
||
**Causes:**
|
||
|
||
1. Model not registered in `registry.py`
|
||
2. Wrong class name in registry
|
||
3. Missing `__init__.py` exports
|
||
|
||
|
||
**Issue: Shape mismatch in attention**
|
||
|
||
**Symptoms:** `RuntimeError: shape mismatch` in attention forward
|
||
|
||
**Cause:** Incorrect tensor reshaping for vLLM-Omni's attention interface
|
||
|
||
**Solution:** Ensure correct shapes:
|
||
|
||
```python
|
||
# vLLM-Omni expects: [batch, seq_len, num_heads, head_dim]
|
||
query = query.view(batch_size, seq_len, self.num_heads, self.head_dim)
|
||
key = key.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim)
|
||
value = value.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim)
|
||
|
||
hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
|
||
|
||
# Reshape back: [batch, seq_len, num_heads, head_dim] → [batch, seq_len, hidden_size]
|
||
hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
|
||
```
|
||
|
||
**Issue: Different outputs compared to Diffusers**
|
||
|
||
**Symptoms:** Generated images look different from Diffusers
|
||
|
||
**Causes:**
|
||
|
||
1. Attention backend differences (FlashAttention vs PyTorch SDPA)
|
||
2. Missing normalization or scaling
|
||
|
||
**4. Issue: Out of memory (OOM)**
|
||
|
||
**Symptoms:** CUDA out of memory errors
|
||
|
||
**Solutions:**
|
||
|
||
1. **Reduce batch size:**
|
||
```python
|
||
omni.generate(prompts=[...], max_num_seqs=2)
|
||
```
|
||
|
||
2. **Use smaller image size:**
|
||
```python
|
||
sampling_params = OmniDiffusionSamplingParams(height=512, width=512)
|
||
```
|
||
|
||
3. **Enable model offloading:**
|
||
```python
|
||
omni = Omni(model="...", enable_cpu_offload=True)
|
||
```
|
||
|
||
4. **Apply vae tiling and slicing**
|
||
```python
|
||
omni = Omni(model="...", vae_use_slicing=True, vae_use_tiling=True,)
|
||
```
|
||
|
||
---
|
||
|
||
## Pull Request Checklist
|
||
|
||
When submitting a PR to add your model, include:
|
||
|
||
**1. Implementation Files**
|
||
|
||
- ✅ Transformer model (`xxx_transformer.py`)
|
||
- ✅ Pipeline (`pipeline_xxx.py`)
|
||
- ✅ Registry entries in `registry.py`
|
||
- ✅ `__init__.py` with proper exports
|
||
|
||
**2. Example and Tests**
|
||
|
||
- ✅ Example script in `examples/`
|
||
- ✅ Test file in `tests/e2e/`
|
||
- ✅ Documentation (`docs/`) creation or updates
|
||
|
||
_Note: End-to-end test files in `tests/e2e/` are optional but strongly recommended. README updates are required for all new models._
|
||
|
||
**3. Documentation Updates**
|
||
|
||
- ✅ Add model to supported models table in `docs/models/supported_models.md`
|
||
- ✅ If supporting acceleration features (e.g., sequence parallelism, CFG parallel), update acceleration feature tables in:
|
||
- `docs/user_guide/diffusion_acceleration.md`
|
||
- `docs/user_guide/diffusion/parallelism_acceleration.md`
|
||
|
||
---
|
||
|
||
## Model Recipe
|
||
|
||
After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni.
|
||
|
||
**What to Include**
|
||
|
||
Your recipe should include:
|
||
|
||
1. **Model Overview**: Brief description of the model and its capabilities
|
||
2. **Installation Instructions**: Step-by-step setup instructions including:
|
||
- Installing vllm-omni and dependencies
|
||
- Installing any additional required packages (e.g., xformers, diffusers)
|
||
- Any version requirements
|
||
3. **Usage Examples**: Command-line examples demonstrating how to run the model
|
||
4. **Configuration Details**: Important configuration parameters and their meanings
|
||
|
||
**Example**
|
||
|
||
For reference, see the [LongCat recipe example](https://github.com/vllm-project/recipes/pull/179) which demonstrates the expected format and structure.
|
||
|
||
**Recipe Location**
|
||
|
||
Create your recipe file in the appropriate directory structure:
|
||
- For organization-specific models: `OrganizationName/ModelName.md`
|
||
- For general models: `ModelName.md`
|
||
|
||
The recipe should be a Markdown file that provides clear, reproducible instructions for users to get started with your model.
|
||
|
||
---
|
||
|
||
## Reference Implementations
|
||
|
||
Study these complete examples:
|
||
|
||
| Model | Architecture | Key Features | Files |
|
||
|-------|--------------|--------------|-------|
|
||
| **Qwen-Image** | Dual-stream transformer | CFG-Parallel, SP, TP, Cache | `vllm_omni/diffusion/models/qwen_image/` |
|
||
| **Wan2.2** | Video transformer | Dual transformers, SP, CFG-Parallel | `vllm_omni/diffusion/models/wan2_2/` |
|
||
|
||
---
|
||
|
||
## Summary
|
||
|
||
Adding a diffusion model to vLLM-Omni involves:
|
||
|
||
1. ✅ **Adapt transformer** - Replace attention, remove mixins, add config support
|
||
2. ✅ **Adapt pipeline** - Change interface, add request handling, extract processing
|
||
3. ✅ **Register model** - Add entries to `registry.py`
|
||
4. ✅ **Add examples** - Provide runnable scripts
|
||
5. ✅ **Test thoroughly** - Verify correctness and performance
|
||
6. ✅ **Add advanced features** - Enable parallelism and acceleration (optional)
|
||
7. ✅ **Submit PR** - Include verification results and documentation
|
||
|
||
**Need help?** Check reference implementations or ask in [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
|