Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed.
### 1.2 Pre-Processing and Post-Processing Extraction
Extract the pre-processing and post-processing logic from the pipeline class to follow vLLM-Omni’s execution flow. For Qwen-Image-Edit:
```python
defget_qwen_image_edit_pre_process_func(
od_config:OmniDiffusionConfig,
):
"""
Define a pre-processing function that resizes input images and
pre-process for subsequent inference.
"""
```
```python
defget_qwen_image_edit_post_process_func(
od_config:OmniDiffusionConfig,
):
"""
Defines a post-processing function that post-process images.
"""
```
### 1.3 Define the forward function
The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are:
+ As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly.
### 1.4 Replace some ops or layers in DiT component
vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers.
Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations:
+ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE.
+ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism.
+ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed.
### 1.5 Provide a `_repeated_blocks` in DiT model
`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer.
It's used for torch compile optimizations.
```python
_repeated_blocks=["QwenImageTransformerBlock"]
```
### 1.6 (Optional) implement sequence parallelism
vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic.
You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md)
### 1.7 (Optional) integrate with Cache-Dit
vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py`
User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model.
For each newly integrated model, we need to provide examples script under the examples/ to demonstrate how to initialize the pipeline with Omni, pass in user inputs, and generate outputs.
Key point for writing the example:
+ Use the Omni entrypoint to load the model and construct the pipeline.
+ Show how to format user inputs and pass them via omni.generate(...).
+ Demonstrate the common runtime arguments, such as:
If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance.
### 5.1 Add Extractor Function
First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks:
- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps.
- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity.
- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data.
**Troubleshooting:**
- If results seem unreliable, try:
- Increasing number of prompts (100+ recommended)
- Using more diverse prompts from multiple datasets
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.
