Add Diffusion Model

Adding a Diffusion Model to vLLM-Omni

Overview

This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.

Prerequisites

Before starting, determine:

1. Model category: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
2. Reference source: Diffusers pipeline, custom repo, or a combination
3. Model HuggingFace ID or local checkpoint path
4. Architecture: Scheduler, text encoder, VAE, transformer/backbone

Step 0: Classify the Migration Path

Check the model's HF repo for model_index.json. This determines your path:

Before coding, write a short integration contract covering runtime ownership, checkpoint discovery, reference revision, I/O geometry, attention semantics, CFG behavior, auxiliary components, target hardware, and the default deployment. For Path B, hybrid models, or any optimization work, read references/native-model-integration-checklist.md and use its phase gates.

Path A: Diffusers-Based Model

For models with a standard diffusers layout. See references/transformer-adaptation.md for detailed code patterns.

A1. Analyze model_index.json

Identify components: transformer, scheduler, vae, text_encoder, tokenizer.

A2. Create model directory

vllm_omni/diffusion/models/your_model_name/
├── __init__.py
├── pipeline_your_model.py
└── your_model_transformer.py

A3. Adapt transformer

1. Copy from diffusers source. Remove mixins (ModelMixin, ConfigMixin, AttentionModuleMixin).
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 for acceleration features such as _repeated_blocks and _layerwise_offload_blocks_attrs (see references/transformer-adaptation.md for examples).

A4. Adapt pipeline

Inherit from nn.Module. The key contract:

class YourPipeline(nn.Module):
    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
        # Load VAE, text encoder, tokenizer via from_pretrained()
        # Instantiate transformer (weights loaded later via weights_sources)
        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=od_config.model, subfolder="transformer",
                prefix="transformer.", fall_back_to_pt=True)]

    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
        # Encode prompt → prepare latents → denoise loop → VAE decode
        return DiffusionOutput(output=output)

    def load_weights(self, weights):
        return AutoWeightsLoader(self).load_weights(weights)

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.

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: Native / Non-Diffusers Model

For models without a Diffusers pipeline—weights in custom formats and model code in another public or private repository. Treat that repository as a pinned correctness oracle. A request for native support means the vLLM-Omni runtime must not import the reference implementation; an external adapter is appropriate only when the requested scope explicitly permits that dependency. See references/custom-model-patterns.md for concrete integration patterns.

B1. Understand the reference repo

Study the original model's code to identify:

• Model architecture files (transformers, fusion modules, embeddings)
• Weight format (safetensors, .pth, custom checkpoint structure)
• Weight loading helpers (custom init functions, checkpoint loaders)
• Pre/post-processing (image/audio transforms, tokenization, VAE encode/decode)
• External dependencies (packages not on PyPI)
• Config format (JSON config files, hardcoded dicts)

B2. Decide what lives WHERE

This is the key design decision for custom models. Follow these placement rules:

B3. Handle external dependencies

If the model's code lives in a separate git repo, first decide whether the requested deliverable is native support or an external adapter. Do not silently choose the adapter path.

Option 1: Port the code directly (default for native support)

Copy the essential model files into vllm_omni/diffusion/models/<name>/ and adapt them to shared vLLM-Omni contracts. Keep checkpoint loading strict and use the pinned reference only to generate parity evidence.

Option 2: Import with graceful fallback (adapter scope only)

try:
    from external_model.utils import init_vae, load_checkpoint
except ImportError:
    raise ImportError(
        "Failed to import from dependency 'external_model'. "
        "Please run the download script first."
    )

Use an external runtime dependency only when the user explicitly requested an adapter or a maintainer approved the exception. Document the dependency, pinned revision, installation path, and unsupported native features.

B4. Handle custom weight loading

Custom models have two common patterns for weight loading:

Pattern 1: Bypass standard loader (eager custom init)

When the original model has complex custom init functions that load weights in __init__:

class CustomPipeline(nn.Module):
    def __init__(self, *, od_config, prefix=""):
        super().__init__()
        model = od_config.model
        # Load everything eagerly in __init__ using custom helpers
        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"),
        )
        # NO weights_sources defined — bypasses standard loader

    def load_weights(self, weights):
        pass  # No-op — all weights loaded in __init__

Pattern 2: Use standard loader with custom load_weights (BAGEL style)

When weights are in safetensors format but need name remapping:

from vllm.model_executor.model_loader.weight_utils import default_weight_loader

class CustomPipeline(nn.Module):
    def __init__(self, *, od_config, prefix=""):
        super().__init__()
        # Instantiate model architecture without weights
        self.bagel = BagelModel(config)
        self.vae = AutoEncoder(ae_params)

        # Point loader at the safetensors in the model root
        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=od_config.model,
                subfolder=None,  # weights at root, not in subfolder
                prefix="",
                fall_back_to_pt=False,
            )
        ]

    def load_weights(self, weights):
        # Custom name remapping for non-diffusers weight names
        params = dict(self.named_parameters())
        loaded = set()
        for name, tensor in weights:
            # Remap original weight names to vllm-omni module names
            name = self._remap_weight_name(name)
            if name in params:
                param = params[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, tensor)
                loaded.add(name)
        return loaded

B5. Create the model_index.json

Prefer a model_index.json at the model root when vLLM-Omni owns an assembled checkpoint directory. For custom models, this is minimal:

{
    "_class_name": "YourModelPipeline",
    "custom_key": "path/to/custom_weights.safetensors"
}

The _class_name must match a key in _DIFFUSION_MODELS in registry.py. Additional keys are model-specific (accessed via od_config.model_config).

If the released repository is immutable and has neither a root config.json nor a Diffusers index, add it to a generic native-checkpoint signature resolver: match the exact Hub ID or a distinctive set of local files and return the pipeline class. Do not use model-name substrings or add parallel one-off predicates in CLI, config, and serving consumers.

If the model's weights come from multiple HF repos, write a download script that:

1. Downloads from each repo
2. Assembles into a single directory
3. Generates model_index.json
4. Installs any external dependencies (git clone + .pth file)

Place at: examples/offline_inference/<name>/download_<name>.py

B6. Handle multi-modal inputs

If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from vllm_omni/diffusion/models/interface.py:

from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput

class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
    # Protocol markers — the engine uses these to enable proper input routing
    pass

Preprocessing for custom models is typically done inside forward() rather than via registered pre-process functions, since the logic is often tightly coupled to the model.

B7. Continue at Step 4 below.

Common Steps (Both Paths)

Step 4: Register Model in registry.py

Edit vllm_omni/diffusion/registry.py:

_DIFFUSION_MODELS = {
    "YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
}
_DIFFUSION_POST_PROCESS_FUNCS = {
    "YourModelPipeline": "get_your_model_post_process_func",  # if applicable
}
_DIFFUSION_PRE_PROCESS_FUNCS = {
    "YourModelPipeline": "get_your_model_pre_process_func",  # if applicable
}

The registry key is the _class_name from model_index.json. The tuple is (folder_name, module_file, class_name).

Create __init__.py exporting the pipeline class and any factory functions.

Step 5: Run, Test, Debug

Use the appropriate existing example script:

Reuse these shared scripts even for custom models when their request and output contracts fit. Create a dedicated model script only when the shared category cannot represent the protocol, and document that gap in the PR.

Validation: No errors, output is meaningful, quality matches reference implementation.

See references/troubleshooting.md for common errors.

Step 6: Add Example Scripts

Only when the shared category scripts cannot represent the model, create:

• examples/offline_inference/your_model_name/ — offline script + README
• examples/online_serving/your_model_name/ — server script + client
• Download script if weights require assembly from multiple sources

Step 7: Update Documentation

Follow the add-recipe skill to add or update the model-family recipe and its recipes/README.md row with verified specifications, hardware, commands, feature links, and qualification evidence.

Required updates:

1. docs/user_guide/diffusion/parallelism/overview.md — parallelism support overview/table
2. docs/user_guide/diffusion/cpu_offload.md — if CPU offload supported (add to supported models table)
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

Follow the vllm-omni-test skill for markers, file naming, Buildkite wiring, and run commands. Also read l4_functionality_tests.inc.md, test_system_overview.md, and test_writing_guide.md.

Classify the model's CI priority first:

Per-level deliverables (diffusion / pytest.mark.diffusion):

L2 & L3 online — same file, dual marks on the baseline smoke: The first / simplest case in test_{slug}.py (default deploy, minimal steps, single prompt) should carry both @pytest.mark.core_model and @pytest.mark.advanced_model on the same function so L2 (test-ready.yml) and L3 (test-merge.yml) share one smoke test. Heavier deploy variants or API paths in the same file use advanced_model only. When L3 moves to nightly, migrate those heavier cases into test_{slug}_expansion.py with full_model and remove the dedicated test-merge.yml job (see test_longcat_image_expansion.py, test_qwen_image_expansion.py).

L4 design (high priority): Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into few parametrized OmniServerParams rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (cards_1 vs not cards_1).

L4 design (medium / low): One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority.

Reference implementations: tests/e2e/online_serving/test_qwen_image_edit_expansion.py, tests/e2e/online_serving/test_longcat_image_expansion.py, tests/e2e/online_serving/test_hunyuan_video_15_expansion.py.

Keep model-specific code inside test modules — not tests/helpers/{slug}.py: deploy constants, prompts, sampling dicts, and inline request_config / form_data belong in each test_{slug}.py and test_{slug}_expansion.py. Do not add per-model files under tests/helpers/; reuse only repo-wide harness (mark, media, runtime, stage_config, assertions). L2+ online/offline e2e: reuse or add send_*_request in tests/helpers/runtime.py — tests call the handler, not raw omni.generate / HTTP. See vllm-omni-test skill § Runtime send helpers.

Keep the model suite proportional using the six distinct failure owners in the native integration checklist. Combine supported features into a few parametrized E2E rows instead of creating one test per optimization.