# Distilled Stable Diffusion inference

## Distilled Stable Diffusion inference

![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)![Open In Studio Lab](https://studiolab.sagemaker.aws/studiolab.svg)

Stable Diffusion inference can be a computationally intensive process because it must iteratively denoise the latents to generate an image. To reduce the computational burden, you can use a _distilled_ version of the Stable Diffusion model from [Nota AI](https://huggingface.co/nota-ai). The distilled version of their Stable Diffusion model eliminates some of the residual and attention blocks from the UNet, reducing the model size by 51% and improving latency on CPU/GPU by 43%.

Read this [blog post](https://huggingface.co/blog/sd_distillation) to learn more about how knowledge distillation training works to produce a faster, smaller, and cheaper generative model.

Let’s load the distilled Stable Diffusion model and compare it against the original Stable Diffusion model:

Copied

```
from diffusers import StableDiffusionPipeline
import torch

distilled = StableDiffusionPipeline.from_pretrained(
    "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")

original = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")
```

Given a prompt, get the inference time for the original model:

Copied

```
import time

seed = 2023
generator = torch.manual_seed(seed)

NUM_ITERS_TO_RUN = 3
NUM_INFERENCE_STEPS = 25
NUM_IMAGES_PER_PROMPT = 4

prompt = "a golden vase with different flowers"

start = time.time_ns()
for _ in range(NUM_ITERS_TO_RUN):
    images = original(
        prompt,
        num_inference_steps=NUM_INFERENCE_STEPS,
        generator=generator,
        num_images_per_prompt=NUM_IMAGES_PER_PROMPT
    ).images
end = time.time_ns()
original_sd = f"{(end - start) / 1e6:.1f}"

print(f"Execution time -- {original_sd} ms\n")
"Execution time -- 45781.5 ms"
```

Time the distilled model inference:

Copied

```
start = time.time_ns()
for _ in range(NUM_ITERS_TO_RUN):
    images = distilled(
        prompt,
        num_inference_steps=NUM_INFERENCE_STEPS,
        generator=generator,
        num_images_per_prompt=NUM_IMAGES_PER_PROMPT
    ).images
end = time.time_ns()

distilled_sd = f"{(end - start) / 1e6:.1f}"
print(f"Execution time -- {distilled_sd} ms\n")
"Execution time -- 29884.2 ms"
```

<figure><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/original_sd.png" alt=""><figcaption><p>original Stable Diffusion (45781.5 ms)</p></figcaption></figure>

<figure><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/distilled_sd.png" alt=""><figcaption><p>distilled Stable Diffusion (29884.2 ms)</p></figcaption></figure>

### Tiny AutoEncoder

To speed inference up even more, use a tiny distilled version of the [Stable Diffusion VAE](https://huggingface.co/sayakpaul/taesdxl-diffusers) to denoise the latents into images. Replace the VAE in the distilled Stable Diffusion model with the tiny VAE:

Copied

```
from diffusers import AutoencoderTiny

distilled.vae = AutoencoderTiny.from_pretrained(
    "sayakpaul/taesd-diffusers", torch_dtype=torch.float16, use_safetensors=True,
).to("cuda")
```

Time the distilled model and distilled VAE inference:

Copied

```
start = time.time_ns()
for _ in range(NUM_ITERS_TO_RUN):
    images = distilled(
        prompt,
        num_inference_steps=NUM_INFERENCE_STEPS,
        generator=generator,
        num_images_per_prompt=NUM_IMAGES_PER_PROMPT
    ).images
end = time.time_ns()

distilled_tiny_sd = f"{(end - start) / 1e6:.1f}"
print(f"Execution time -- {distilled_tiny_sd} ms\n")
"Execution time -- 27165.7 ms"
```

<figure><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/distilled_sd_vae.png" alt=""><figcaption><p>distilled Stable Diffusion + Tiny AutoEncoder (27165.7 ms)</p></figcaption></figure>
