274 lines
9.6 KiB
Markdown
274 lines
9.6 KiB
Markdown
# Custom Model Patterns Reference
|
|
|
|
Patterns for adding models that don't come from the standard diffusers pipeline format.
|
|
|
|
## Directory Structure Comparison
|
|
|
|
### Diffusers-based model (e.g., Wan2.2)
|
|
|
|
```
|
|
vllm_omni/diffusion/models/wan2_2/
|
|
├── __init__.py # Exports pipeline + transformer + helpers
|
|
├── pipeline_wan2_2.py # Pipeline: loads components via from_pretrained()
|
|
├── pipeline_wan2_2_i2v.py # Variant pipeline for image-to-video
|
|
└── wan2_2_transformer.py # Transformer: ported from diffusers, uses Attention layer
|
|
```
|
|
|
|
The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
|
|
|
|
### Custom model with external deps (e.g., DreamID-Omni)
|
|
|
|
```
|
|
vllm_omni/diffusion/models/dreamid_omni/
|
|
├── __init__.py # Exports pipeline only
|
|
├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
|
├── fusion.py # Custom fusion architecture (video + audio cross-attention)
|
|
└── wan2_2.py # Re-implemented Wan backbone with split API
|
|
|
|
examples/offline_inference/x_to_video_audio/
|
|
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo
|
|
```
|
|
|
|
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
|
|
|
|
### Custom model with ported code (e.g., BAGEL)
|
|
|
|
```
|
|
vllm_omni/diffusion/models/bagel/
|
|
├── __init__.py
|
|
├── pipeline_bagel.py # Pipeline: instantiates models, uses weights_sources
|
|
├── bagel_transformer.py # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
|
|
└── autoencoder.py # Custom VAE ported from original repo
|
|
```
|
|
|
|
Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
|
|
|
|
## Weight Loading Patterns
|
|
|
|
### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
|
|
|
|
```
|
|
init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
|
|
```
|
|
|
|
- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
|
|
- `load_weights()` receives `(name, tensor)` pairs from the loader
|
|
- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
|
|
|
|
### Pattern 2: Custom safetensors at root (BAGEL)
|
|
|
|
```
|
|
init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
|
|
```
|
|
|
|
- `weights_sources` points to **root** of model directory, not a subfolder
|
|
- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
|
|
- `load_weights()` does heavy name normalization
|
|
|
|
```python
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder=None, # root directory
|
|
prefix="", # no prefix stripping
|
|
fall_back_to_pt=False,
|
|
)
|
|
]
|
|
```
|
|
|
|
### Pattern 3: Fully custom loading (DreamID-Omni)
|
|
|
|
```
|
|
init → load ALL weights eagerly via custom helpers → load_weights() = no-op
|
|
```
|
|
|
|
- No `weights_sources` attribute — standard loader finds nothing to iterate
|
|
- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
|
|
- `load_weights()` is `pass`
|
|
- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
|
|
|
|
Use this when:
|
|
- The original model has complex, well-tested loading code you don't want to rewrite
|
|
- Weights span multiple HF repos
|
|
- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
|
|
|
|
## model_index.json for Custom Models
|
|
|
|
Standard diffusers `model_index.json`:
|
|
```json
|
|
{
|
|
"_class_name": "WanPipeline",
|
|
"_diffusers_version": "0.35.0.dev0",
|
|
"scheduler": ["diffusers", "UniPCMultistepScheduler"],
|
|
"transformer": ["diffusers", "WanTransformer3DModel"],
|
|
"vae": ["diffusers", "AutoencoderKLWan"]
|
|
}
|
|
```
|
|
|
|
Custom model `model_index.json` (minimal):
|
|
```json
|
|
{
|
|
"_class_name": "DreamIDOmniPipeline",
|
|
"fusion": "DreamID-Omni/dreamid_omni.safetensors"
|
|
}
|
|
```
|
|
|
|
The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
|
|
|
|
## External Dependency Management
|
|
|
|
### Git clone + .pth injection (DreamID-Omni pattern)
|
|
|
|
```python
|
|
def download_dependency():
|
|
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
with open(LOCK_FILE, "w") as f:
|
|
fcntl.flock(f, fcntl.LOCK_EX)
|
|
if not DEPENDENCY_DIR.exists():
|
|
subprocess.run([
|
|
"git", "clone", "--depth", "1",
|
|
REPO_URL, "--branch", BRANCH,
|
|
str(DEPENDENCY_DIR)
|
|
], check=True)
|
|
fcntl.flock(f, fcntl.LOCK_UN)
|
|
|
|
# Add to Python path via .pth file
|
|
site_packages = Path(site.getsitepackages()[0])
|
|
pth_file = site_packages / "vllm_omni_dependency.pth"
|
|
pth_file.write_text(str(DEPENDENCY_DIR))
|
|
```
|
|
|
|
### Direct port (BAGEL pattern)
|
|
|
|
Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
|
|
|
|
## Multi-Modal Input/Output Protocols
|
|
|
|
Custom models that handle images, audio, or video I/O should implement protocol classes:
|
|
|
|
```python
|
|
from vllm_omni.diffusion.models.interface import (
|
|
SupportImageInput, # Model accepts image input
|
|
SupportAudioInput, # Model accepts audio input
|
|
SupportAudioOutput, # Model produces audio output
|
|
)
|
|
|
|
class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
|
|
pass # Protocol markers enable proper engine routing
|
|
```
|
|
|
|
The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
|
|
|
|
## Hardcoded Config vs Config Files
|
|
|
|
Diffusers models use `config.json` in each subfolder. Custom models often use:
|
|
|
|
**Module-level config dicts** (DreamID-Omni):
|
|
```python
|
|
VIDEO_CONFIG = {
|
|
"patch_size": [1, 2, 2], "model_type": "ti2v",
|
|
"dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
|
|
}
|
|
```
|
|
|
|
**Loaded from custom JSON** (BAGEL):
|
|
```python
|
|
cfg_path = os.path.join(model_path, "config.json")
|
|
with open(cfg_path) as f:
|
|
bagel_cfg = json.load(f)
|
|
vae_cfg = bagel_cfg.get("vae_config", {})
|
|
```
|
|
|
|
## Custom Architecture Patterns
|
|
|
|
### Split forward API (DreamID-Omni)
|
|
|
|
When a fusion model needs to interleave blocks from two backbones:
|
|
|
|
```python
|
|
class WanModel(nn.Module):
|
|
def prepare_transformer_block_kwargs(self, x, t, context, ...):
|
|
# Patch embed, time embed, text embed, RoPE
|
|
return x, e, kwargs
|
|
|
|
def post_transformer_block_out(self, x, grid_sizes, e):
|
|
# Output projection, unpatchify
|
|
return output
|
|
|
|
def forward(self, *args, **kwargs):
|
|
raise NotImplementedError # Fusion model handles block iteration
|
|
```
|
|
|
|
The `FusionModel` then iterates blocks in lock-step:
|
|
```python
|
|
for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
|
|
video_out = video_block(video_hidden, ...)
|
|
audio_out = audio_block(audio_hidden, ...)
|
|
# Cross-attend between modalities
|
|
video_out = cross_attention(video_out, audio_out)
|
|
audio_out = cross_attention(audio_out, video_out)
|
|
```
|
|
|
|
### LLM-as-denoiser (BAGEL)
|
|
|
|
When the backbone is a language model that also does diffusion:
|
|
|
|
```python
|
|
class BagelModel(nn.Module):
|
|
def __init__(self):
|
|
self.language_model = Qwen2MoTForCausalLM(config)
|
|
self.vit_model = SiglipVisionModel(vit_config)
|
|
```
|
|
|
|
The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
|
|
|
|
## Pre/Post Processing for Custom Models
|
|
|
|
Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
|
|
|
|
```python
|
|
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
|
# Inline preprocessing
|
|
image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
|
|
image_latent = self._vae_encode(image)
|
|
|
|
# ... denoising loop ...
|
|
|
|
# Inline postprocessing
|
|
pil_image = self._decode_to_pil(latents)
|
|
return DiffusionOutput(output=[pil_image])
|
|
```
|
|
|
|
If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
|
|
|
|
## Download Script Template
|
|
|
|
```python
|
|
# examples/offline_inference/<name>/download_<name>.py
|
|
from huggingface_hub import snapshot_download
|
|
import json, os
|
|
|
|
def main(output_dir):
|
|
# Download model weights from HF
|
|
snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
|
|
|
|
# Download additional components if from separate repos
|
|
snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
|
|
allow_patterns=["*.safetensors"])
|
|
|
|
# Generate model_index.json
|
|
config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
|
|
with open(os.path.join(output_dir, "model_index.json"), "w") as f:
|
|
json.dump(config, f, indent=2)
|
|
|
|
# Install external code dependency (if needed)
|
|
download_dependency()
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--output-dir", default="./your_model")
|
|
args = parser.parse_args()
|
|
main(args.output_dir)
|
|
```
|