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

> Learn to implement custom samplers like DDIM, PLMS, or DPM-Solver in Stable Diffusion. Discover how the lightweight interface allows easy integration of new sampling algorithms.

- Repository: [CompVis - Computer Vision and Learning LMU Munich/stable-diffusion](https://github.com/CompVis/stable-diffusion)
- Tags: how-to-guide
- Published: 2026-03-01

---

**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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py), while [`scripts/sample_diffusion.py`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddim.py).

Create a new file in `ldm/models/diffusion/` (e.g., [`my_sampler.py`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/sample_diffusion.py).

```python

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

# 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`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/sample_diffusion.py):

```python

# 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`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/my_custom_sampler.py).

```python
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`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/sample_diffusion.py) to import PLMS:

```python
from ldm.models.diffusion.plms import PLMSSampler

```

Then run:

```bash
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`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/sample_diffusion.py):

```python
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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/my_sampler.py)) to maintain consistency with the existing codebase. Import your class in [`scripts/sample_diffusion.py`](https://github.com/CompVis/stable-diffusion/blob/main/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.