How to Implement Custom Samplers Like DDIM, PLMS, or DPM-Solver in Stable Diffusion

Stable Diffusion decouples the latent diffusion UNet from sampling algorithms through a lightweight Python interface; any class implementing a sample() method that accepts a model and returns latents can serve as a custom sampler, with DDIM, PLMS, and DPM-Solver provided as reference implementations in the CompVis/stable-diffusion repository.

The CompVis/stable-diffusion repository organizes inference into modular components, allowing researchers to implement custom samplers like DDIM, PLMS, or DPM-Solver in Stable Diffusion without modifying the underlying UNet architecture. By adhering to a simple contract—receiving a trained model and implementing a sample method that returns latent tensors—you can integrate deterministic, multistep, or ODE-based sampling strategies. This guide examines the source code structure and provides runnable implementations based on the actual repository files.

Core Architecture

Stable Diffusion separates model definition (the latent diffusion UNet) from sampling algorithms. A sampler is a thin Python class that receives a trained diffusion model and implements a sample method returning latents (and optionally intermediate states).

The repository provides three canonical sampler implementations:

• DDIMSampler in ldm/models/diffusion/ddim.py — Implements deterministic and stochastic DDIM sampling with its own schedule management (make_schedule) and a ddim_sampling loop.
• PLMSSampler in ldm/models/diffusion/plms.py — Uses a pseudo-linear-multistep (Adams-Bashforth) approach, reusing the DDIM schedule but adding a multistep predictor-corrector.
• DPMSolverSampler in ldm/models/diffusion/dpm_solver/sampler.py — Wraps the official DPM-Solver ODE implementation for high-order Runge-Kutta-like steps.

Utility functions for building timesteps, computing alphas, and generating noise reside in ldm/modules/diffusionmodules/util.py, while scripts/sample_diffusion.py demonstrates end-to-end inference.

How a Sampler Works

All samplers follow a consistent four-phase pattern as implemented in the source code:

1. Constructor — Stores the diffusion model and registers constant tensors (betas, alphas_cumprod) on the GPU using register_buffer.

2. Schedule Creation — The make_schedule method builds ddim_timesteps and pre-computes parameters (ddim_alphas, ddim_sigmas) from the model's beta schedule via utilities in ldm.modules.diffusionmodules.util.

3. Sampling Loop — The sample method iterates over timesteps in reverse order:

   • Constructs a timestep tensor ts for the batch.
   • Calls model.apply_model(x, ts, conditioning) to obtain predicted noise e_t.
   • Computes pred_x0 (predicted clean latent) and x_prev (next latent) using pre-computed coefficients.
   • Injects optional classifier-free guidance via unconditional_guidance_scale.

4. Return — Outputs a final latent tensor (and optional intermediates), which the caller decodes using model.decode_first_stage.

The samplers differ only in the predictor step:

• DDIM uses a single-step update.
• PLMS computes an Adams-Bashforth update (e_t_prime) from up to four previous noise predictions.
• DPM-Solver delegates to an external ODE solver with high-order steps.

Adding a Custom Sampler

To implement a new sampling algorithm, create a Python class that adheres to the interface established in ldm/models/diffusion/ddim.py.

Create a new file in ldm/models/diffusion/ (e.g., my_sampler.py) and implement:

• __init__(self, model, **kwargs) — Store the model and register constants using self.register_buffer.
• make_schedule(self, num_steps, ...) — Reuse make_ddim_timesteps and make_ddim_sampling_parameters if basing your sampler on DDIM timesteps, or define your own schedule.
• sample(self, S, batch_size, shape, conditioning=None, ...) — Orchestrate the sampling loop.
• *_sampling(self, ...) — The core loop implementation; replace p_sample_ddim or p_sample_plms with your custom predictor (e.g., p_sample_my).

Return the latent tensor and optional intermediates exactly as the built-in samplers do. Because inference scripts expect only a sample method, you can swap implementations without further modifications.

Using a Sampler in Practice

Below is a minimal, self-contained example demonstrating how to load a checkpoint, instantiate a sampler, and generate images. This mirrors the usage in scripts/sample_diffusion.py.

# -------------------------------------------------

# 1️⃣ Load the checkpoint & config

# -------------------------------------------------

from omegaconf import OmegaConf
from ldm.util import instantiate_from_config
import torch
import yaml

def load_model_from_config(config, ckpt_path):
    model = instantiate_from_config(config)
    sd = torch.load(ckpt_path, map_location="cpu")["state_dict"]
    model.load_state_dict(sd, strict=False)
    model.cuda()
    model.eval()
    return model

config_path = "logs/config.yaml"
ckpt_path = "logs/model.ckpt"

cfg = OmegaConf.load(config_path)
model = load_model_from_config(cfg.model, ckpt_path)

# -------------------------------------------------

# 2️⃣ Choose a sampler

# -------------------------------------------------

from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
from ldm.models.diffusion.dpm_solver.sampler import DPMSolverSampler

# sampler = DDIMSampler(model)       # deterministic/stochastic DDIM

# sampler = PLMSSampler(model)       # multistep PLMS

sampler = DPMSolverSampler(model)    # high-order DPM-Solver

# -------------------------------------------------

# 3️⃣ Define sampling parameters

# -------------------------------------------------

steps = 50
batch = 4
shape = (batch, model.model.diffusion_model.in_channels,
         model.model.diffusion_model.image_size,
         model.model.diffusion_model.image_size)

# -------------------------------------------------

# 4️⃣ Run the sampler

# -------------------------------------------------

samples, _ = sampler.sample(
    S=steps,
    batch_size=batch,
    shape=shape,
    conditioning=None,                 # Replace with text embeddings for txt2img

    eta=0.0,                           # DDIM-specific; ignored by PLMS/DPM-Solver

    unconditional_guidance_scale=7.5,  # Classifier-free guidance

)

# -------------------------------------------------

# 5️⃣ Decode latents to RGB

# -------------------------------------------------

decoded = model.decode_first_stage(samples)
decoded = (decoded.clamp(-1, 1) + 1) / 2
decoded = decoded.cpu().permute(0, 2, 3, 1).numpy()

Replace conditioning with a token tensor from the CLIP text encoder (model.cond_stage_model) for text-to-image generation.

To switch samplers in the official script, modify the import in scripts/sample_diffusion.py:

# from ldm.models.diffusion.ddim import DDIMSampler

from ldm.models.diffusion.plms import PLMSSampler
sampler = PLMSSampler(model)

No other code changes are required because the script calls sampler.sample(...) with the same signature.

Code Examples

Minimal Custom Sampler

This example copies the DDIM skeleton and implements a simple Euler predictor. Save as ldm/models/diffusion/my_custom_sampler.py.

import torch
import numpy as np
from tqdm import tqdm
from ldm.modules.diffusionmodules.util import (
    make_ddim_timesteps, make_ddim_sampling_parameters, noise_like
)

class MyCustomSampler:
    def __init__(self, model):
        self.model = model
        self.ddpm_num_timesteps = model.num_timesteps

    def register_buffer(self, name, attr):
        if isinstance(attr, torch.Tensor) and attr.device != torch.device("cuda"):
            attr = attr.to(torch.device("cuda"))
        setattr(self, name, attr)

    def make_schedule(self, steps, discretize="uniform", eta=0.0):
        self.ddim_timesteps = make_ddim_timesteps(
            ddim_discr_method=discretize,
            num_ddim_timesteps=steps,
            num_ddpm_timesteps=self.ddpm_num_timesteps,
        )
        alphas = self.model.alphas_cumprod.cpu()
        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(
            alphacums=alphas,
            ddim_timesteps=self.ddim_timesteps,
            eta=eta,
        )
        self.register_buffer("ddim_sigmas", ddim_sigmas)
        self.register_buffer("ddim_alphas", ddim_alphas)
        self.register_buffer("ddim_alphas_prev", ddim_alphas_prev)

    @torch.no_grad()
    def sample(self, S, batch_size, shape, conditioning=None, eta=0.0):
        self.make_schedule(S, eta=eta)
        C, H, W = shape
        size = (batch_size, C, H, W)
        img = torch.randn(size, device=self.model.betas.device)

        timesteps = np.flip(self.ddim_timesteps)
        for i, step in enumerate(tqdm(timesteps, desc="MyCustomSampler")):
            ts = torch.full((batch_size,), step, device=img.device, dtype=torch.long)
            eps = self.model.apply_model(img, ts, conditioning)
            
            # Custom Euler predictor step

            a_t = self.ddim_alphas[i]
            a_prev = self.ddim_alphas_prev[i]
            sigma = self.ddim_sigmas[i]
            pred_x0 = (img - (1 - a_t).sqrt() * eps) / a_t.sqrt()
            dir_xt = (1 - a_prev - sigma**2).sqrt() * eps
            noise = sigma * noise_like(img.shape, img.device)
            img = a_prev.sqrt() * pred_x0 + dir_xt + noise
            
        return img, None

Using PLMS from Command Line

Edit scripts/sample_diffusion.py to import PLMS:

from ldm.models.diffusion.plms import PLMSSampler

Then run:

python scripts/sample_diffusion.py \
    -r logs/checkpoints/model.ckpt \
    --custom_steps 50 \
    --eta 0.0

Switching to DPM-Solver

Modify the import in scripts/sample_diffusion.py:

from ldm.models.diffusion.dpm_solver.sampler import DPMSolverSampler

The eta parameter is ignored when using DPM-Solver, as this sampler does not implement the DDIM noise parameter.

Summary

• Sampler Interface — Any class implementing sample(S, batch_size, shape, ...) and storing the model in self.model can function as a custom sampler in Stable Diffusion.
• Reference Implementations — DDIMSampler, PLMSSampler, and DPMSolverSampler in ldm/models/diffusion/ demonstrate deterministic, multistep, and ODE-based approaches respectively.
• Schedule Utilities — Reuse make_ddim_timesteps and make_ddim_sampling_parameters from ldm/modules/diffusionmodules/util.py to handle beta schedules and coefficient pre-computation.
• Model Interaction — Always call model.apply_model(x, ts, conditioning) within the sampling loop to obtain noise predictions, then apply your custom update rule.
• Integration — Samplers are plug-and-play; changing the import line in scripts/sample_diffusion.py is sufficient to switch between DDIM, PLMS, DPM-Solver, or custom implementations.

Frequently Asked Questions

What is the difference between DDIM and PLMS samplers in Stable Diffusion?

DDIM performs single-step updates using the DDIM formula, making it suitable for deterministic or lightly stochastic generation with the eta parameter. PLMS (Pseudo Linear Multistep) implements an Adams-Bashforth predictor that uses up to four previous noise estimates (e_t) to compute e_t_prime, often producing smoother transitions with fewer steps but requiring multistep initialization.

How do I register custom diffusion schedules in a new sampler class?

Override the make_schedule method and use self.register_buffer(name, tensor) to store schedule-dependent tensors (like ddim_alphas or custom coefficients) on the GPU. This ensures they move with the model when calling .cuda() and persist across sampling steps without being treated as model parameters.

Can I use the eta (η) parameter with DPM-Solver in Stable Diffusion?

No. The DPM-Solver implementation in ldm/models/diffusion/dpm_solver/sampler.py ignores the eta argument because it is an ODE solver designed for deterministic sampling. The eta parameter only affects stochasticity in DDIM-based samplers; for DPM-Solver, randomness is controlled through the initial noise tensor passed to sample().

Where should I place my custom sampler file in the repository structure?

Create new sampler files in ldm/models/diffusion/ (e.g., ldm/models/diffusion/my_sampler.py) to maintain consistency with the existing codebase. Import your class in scripts/sample_diffusion.py or your inference script using from ldm.models.diffusion.my_sampler import MySampler, ensuring the module path matches the file location.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}