[chore][skills]: tightens the add-diffusion-model skill under .claude/skills to better match the current vllm-omni codebase and repo conventions (#4110)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
LancerandHongsheng Liu authored and GitHub committed 2026-06-22 10:17:27 +08:00
1 parent f0d7fe18f2
commit 36bd8029f3
3 files changed
+43 -15

No files matched your search

+40 -13
View File
@@ -24,7 +24,7 @@ Check the model's HF repo for `model_index.json`. This determines your path:
| Scenario | How to identify | Migration path |
|----------|----------------|----------------|
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed |
| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
@@ -52,7 +52,7 @@ vllm_omni/diffusion/models/your_model_name/
2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
5. Add class attributes for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples).
### A4. Adapt pipeline
@@ -78,13 +78,34 @@ class YourPipeline(nn.Module):
Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
### A4.1 Add progress bar support (recommended)
For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging.
```python
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
class YourPipeline(nn.Module, ProgressBarMixin):
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
# ... prepare timesteps / latents ...
with self.progress_bar(total=len(timesteps)) as progress_bar:
for i, t in enumerate(timesteps):
# predict noise / scheduler step
latents = ...
progress_bar.update()
return DiffusionOutput(output=output)
```
For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`.
### A5. Register, test, docs → continue at Step 4 below.
---
## Path B: Custom/Private Repo Model
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
### B1. Understand the reference repo
@@ -134,7 +155,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
### B4. Handle custom weight loading
Custom models have two patterns for weight loading:
Custom models have two common patterns for weight loading:
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
@@ -149,8 +170,10 @@ class CustomPipeline(nn.Module):
self.vae = custom_init_vae(model, device=self.device)
self.text_encoder = custom_init_text_encoder(model, device=self.device)
self.transformer = CustomFusionModel(CONFIG)
load_custom_checkpoint(self.transformer,
checkpoint_path=os.path.join(model, "model.safetensors"))
load_custom_checkpoint(
self.transformer,
checkpoint_path=os.path.join(model, "model.safetensors"),
)
# NO weights_sources defined — bypasses standard loader
def load_weights(self, weights):
@@ -162,6 +185,8 @@ class CustomPipeline(nn.Module):
When weights are in safetensors format but need name remapping:
```python
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
class CustomPipeline(nn.Module):
def __init__(self, *, od_config, prefix=""):
super().__init__()
@@ -187,7 +212,9 @@ class CustomPipeline(nn.Module):
# Remap original weight names to vllm-omni module names
name = self._remap_weight_name(name)
if name in params:
default_weight_loader(params[name], tensor)
param = params[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, tensor)
loaded.add(name)
return loaded
```
@@ -281,12 +308,12 @@ For Omni or custom models, create:
### Step 7: Update Documentation
Required updates:
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
5. `examples/offline_inference/xxx/README.md` — offline example docs
6. `examples/online_serve/xxx/README.md` — online serve docs
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
6. `examples/online_serving/<name>/README.md` — online serving docs
### Step 8: Add E2E Tests
@@ -537,7 +564,7 @@ class YourTransformer(nn.Module):
#### 10e. Update parallelism documentation
After adding parallelism support, update:
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
1. `docs/user_guide/diffusion/parallelism/overview.md` — add your model to the support overview/table
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
### Step 11: Add CPU Offload Support
@@ -569,4 +569,4 @@ Check which parallelism methods a model supports:
| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.
The canonical per-model support overview/table is in `docs/user_guide/diffusion/parallelism/overview.md`.
@@ -219,13 +219,14 @@ def load_weights(self, weights):
### Class-Level Attributes for Features
```python
class MyTransformer(nn.Module):
# torch.compile: list block class names that repeat and can be compiled
_repeated_blocks = ["MyTransformerBlock"]
# CPU offload: attribute name of the nn.ModuleList containing blocks
_layerwise_offload_blocks_attr = "blocks"
_layerwise_offload_blocks_attrs = ["blocks"]
# LoRA: mapping of fused param names to original param names
packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}