# How to Configure EMA Weights for More Stable Inference in Stable Diffusion

> Achieve stable Stable Diffusion inference using EMA weights. Learn to configure model settings and use EMA scope for smoother, lower-variance image generation.

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

---

**Enable EMA weights in Stable Diffusion by setting `use_ema: True` in your configuration file and wrapping inference code with `with model.ema_scope():` to temporarily swap in the shadow parameters for smoother, lower-variance image generation.**

The CompVis/stable-diffusion repository implements Exponential Moving Average (EMA) weight tracking to improve generation stability during inference. When you configure EMA weights for more stable inference in Stable Diffusion, you leverage a shadow copy of model parameters that maintains lower variance than raw training weights, resulting in more consistent image outputs.

## What Are EMA Weights in Stable Diffusion?

Stable Diffusion maintains a **shadow copy** of the model parameters known as **EMA weights**. During training, this shadow copy is updated each step with a decay factor (default **0.9999**), creating a smoothed version of the parameters that accumulates less noise than the raw training weights. At inference time, swapping in these EMA weights often yields **smoother** and **more reliable** generations because the EMA version has lower variance and better generalization than the non-EMA parameters.

## How EMA Is Implemented in the CompVis/stable-diffusion Codebase

### The LitEma Class (ldm/modules/ema.py)

The core EMA implementation resides in [`ldm/modules/ema.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/ema.py) within the **`LitEma`** class. This class initializes buffers for every trainable parameter in the model and updates them during the forward pass using the specified decay factor. Key methods include:

- **`copy_to()`**: Copies the EMA weights into the model parameters, effectively swapping the shadow weights in for inference.
- **`restore()`**: Restores the original non-EMA parameters back to the model after inference completes.

### Integration with the DDPM Model (ldm/models/diffusion/ddpm.py)

The **`DDPM`** class in [`ldm/models/diffusion/ddpm.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py) (lines 88-92) integrates EMA functionality through the **`use_ema`** configuration flag. When `use_ema=True`, the constructor instantiates the shadow copy:

```python
if self.use_ema:
    self.model_ema = LitEma(self.model)

```

This makes the EMA weights available throughout the model lifecycle, from training through checkpoint saving.

### The ema_scope Context Manager

To simplify inference, the DDPM class provides the **`ema_scope`** context manager (implemented around lines 71-85 in [`ldm/models/diffusion/ddpm.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py)). This context temporarily swaps in EMA weights at entry and automatically restores the original parameters upon exit, ensuring clean state management even if exceptions occur during sampling.

## How to Configure EMA Weights for Inference

### Enabling EMA in Configuration Files

To activate EMA weights for inference, modify the model configuration YAML file. In [`configs/stable-diffusion/v1-inference.yaml`](https://github.com/CompVis/stable-diffusion/blob/main/configs/stable-diffusion/v1-inference.yaml), set the **`use_ema`** parameter to `True`:

```yaml
model:
  params:
    # … other parameters …

    use_ema: True          # ← turn on EMA for inference

```

By default, this value is often set to `False` in inference configurations, so explicit activation is required.

### Customizing the Decay Factor

The default decay factor of **0.9999** can be adjusted for different smoothing behaviors. To customize this value, subclass the DDPM model and override the EMA initialization:

```python

# my_custom_ddpm.py

from ldm.models.diffusion.ddpm import DDPM
from ldm.modules.ema import LitEma

class DDPMWithCustomEMA(DDPM):
    def __init__(self, *args, ema_decay=0.9995, **kwargs):
        super().__init__(*args, **kwargs)
        if self.use_ema:
            # Replace the default LitEma instance with a custom decay

            self.model_ema = LitEma(self.model, decay=ema_decay)

# Load the checkpoint as usual, then replace the class:

model = DDPMWithCustomEMA(**config["model"]["params"])

```

Lower decay values (e.g., 0.9995) allow the EMA to adapt more quickly to recent training changes, while higher values (e.g., 0.9999) provide stronger smoothing.

### Using EMA Weights in Sampling Scripts

The official sampling scripts already implement EMA weight swapping. In [`scripts/txt2img.py`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/txt2img.py), the generation loop wraps sampling with the EMA context manager:

```python

# scripts/txt2img.py – the relevant part

with torch.no_grad():
    with precision_scope("cuda"):
        # EMA weights are swapped in for the whole sampling block

        with model.ema_scope():
            for prompts in data:
                # standard sampling code …

                samples_ddim, _ = sampler.sample(...)

```

No additional code changes are required once `use_ema` is enabled in the configuration. The `ema_scope()` context ensures that EMA weights are active during the forward pass and automatically restores the original parameters when sampling completes.

## Summary

- **EMA weights** are shadow copies of model parameters maintained during training with a decay factor (default 0.9999) to reduce variance.
- The **`LitEma`** class in [`ldm/modules/ema.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/ema.py) implements the core logic, while **`DDPM`** in [`ldm/models/diffusion/ddpm.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py) integrates it via the `use_ema` flag.
- Enable EMA for inference by setting **`use_ema: True`** in [`configs/stable-diffusion/v1-inference.yaml`](https://github.com/CompVis/stable-diffusion/blob/main/configs/stable-diffusion/v1-inference.yaml).
- Use the **`ema_scope()`** context manager to temporarily swap EMA weights during sampling, ensuring automatic restoration of original parameters afterward.
- Customize the **decay factor** by subclassing `DDPM` and passing a different value to `LitEma` (e.g., 0.9995 for faster adaptation).

## Frequently Asked Questions

### What is the default EMA decay factor in Stable Diffusion?

The default decay factor is **0.9999**, defined in the `LitEma` class initialization in [`ldm/modules/ema.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/ema.py). This high value ensures that the EMA weights change very slowly, providing strong smoothing over thousands of training steps. You can specify a lower value (such as 0.9995) if you want the EMA to adapt more quickly to recent training updates.

### Do I need to modify sampling scripts to use EMA weights?

No, the official sampling scripts such as [`scripts/txt2img.py`](https://github.com/CompVis/stable-diffusion/blob/main/scripts/txt2img.py) already include the necessary code to use EMA weights. You only need to ensure that `use_ema: True` is set in your configuration file (such as [`configs/stable-diffusion/v1-inference.yaml`](https://github.com/CompVis/stable-diffusion/blob/main/configs/stable-diffusion/v1-inference.yaml)). The scripts wrap the sampling loop with `with model.ema_scope():`, which automatically handles the weight swapping and restoration.

### Can I switch between EMA and non-EMA weights during inference?

Yes, you can switch between EMA and non-EMA weights dynamically during inference by using the `ema_scope()` context manager. When you enter the context (`with model.ema_scope():`), the EMA weights are copied to the model parameters via `copy_to()`. When you exit the context, the original parameters are automatically restored via `restore()`. This allows you to generate images with EMA weights for some batches and non-EMA weights for others within the same script execution.

### Where are EMA weights stored in the checkpoint?

EMA weights are stored as part of the `LitEma` instance state within the `DDPM` model. Specifically, when `use_ema=True`, the `DDPM` class instantiates `self.model_ema = LitEma(self.model)` in [`ldm/models/diffusion/ddpm.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/diffusion/ddpm.py). The `LitEma` class maintains buffers for every trainable parameter, and these buffers are saved and loaded as part of the PyTorch model state dict. When you load a checkpoint that was trained with EMA enabled, the shadow weights are automatically restored to the `model_ema` instance.