vllm-omni/.claude/skills/add-diffusion-model/references/custom-model-patterns.md

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

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:

{
    "_class_name": "WanPipeline",
    "_diffusers_version": "0.35.0.dev0",
    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
    "transformer": ["diffusers", "WanTransformer3DModel"],
    "vae": ["diffusers", "AutoencoderKLWan"]
}

Custom model model_index.json (minimal):

{
    "_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)

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:

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

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

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:

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:

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:

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:

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

# 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)