#
Load schedulers and models

Diffusion pipelines are a collection of interchangeable schedulers and models that can be mixed and matched to tailor a pipeline to a specific use case. The scheduler encapsulates the entire denoising process such as the number of denoising steps and the algorithm for finding the denoised sample. A scheduler is not parameterized or trained so they don’t take very much memory. The model is usually only concerned with the forward pass of going from a noisy input to a less noisy sample.

This guide will show you how to load schedulers and models to customize a pipeline. You’ll use the stable-diffusion-v1-5/stable-diffusion-v1-5 checkpoint throughout this guide, so let’s load it first.

    import torch
    from diffusers import DiffusionPipeline

    pipeline = DiffusionPipeline.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
    ).to("cuda")

You can see what scheduler this pipeline uses with the `pipeline.scheduler` attribute.

    pipeline.scheduler
    PNDMScheduler {
      "_class_name": "PNDMScheduler",
      "_diffusers_version": "0.21.4",
      "beta_end": 0.012,
      "beta_schedule": "scaled_linear",
      "beta_start": 0.00085,
      "clip_sample": false,
      "num_train_timesteps": 1000,
      "set_alpha_to_one": false,
      "skip_prk_steps": true,
      "steps_offset": 1,
      "timestep_spacing": "leading",
      "trained_betas": null
    }

##
Load a scheduler

Schedulers are defined by a configuration file that can be used by a variety of schedulers. Load a scheduler with the SchedulerMixin.from_pretrained() method, and specify the `subfolder` parameter to load the configuration file into the correct subfolder of the pipeline repository.

For example, to load the DDIMScheduler:

    from diffusers import DDIMScheduler, DiffusionPipeline

    ddim = DDIMScheduler.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="scheduler")

Then you can pass the newly loaded scheduler to the pipeline.

    pipeline = DiffusionPipeline.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", scheduler=ddim, torch_dtype=torch.float16, use_safetensors=True
    ).to("cuda")

##
Compare schedulers

Schedulers have their own unique strengths and weaknesses, making it difficult to quantitatively compare which scheduler works best for a pipeline. You typically have to make a trade-off between denoising speed and denoising quality. We recommend trying out different schedulers to find one that works best for your use case. Call the `pipeline.scheduler.compatibles` attribute to see what schedulers are compatible with a pipeline.

Let’s compare the LMSDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, and the DPMSolverMultistepScheduler on the following prompt and seed.

    import torch
    from diffusers import DiffusionPipeline

    pipeline = DiffusionPipeline.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
    ).to("cuda")

    prompt = "A photograph of an astronaut riding a horse on Mars, high resolution, high definition."
    generator = torch.Generator(device="cuda").manual_seed(8)

To change the pipelines scheduler, use the from_config() method to load a different scheduler’s `pipeline.scheduler.config` into the pipeline.

LMSDiscreteScheduler
EulerDiscreteScheduler
EulerAncestralDiscreteScheduler
DPMSolverMultistepScheduler

LMSDiscreteScheduler typically generates higher quality images than the default scheduler.

    from diffusers import LMSDiscreteScheduler

    pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
    image = pipeline(prompt, generator=generator).images[0]
    image

LMSDiscreteScheduler
EulerDiscreteScheduler
EulerAncestralDiscreteScheduler
DPMSolverMultistepScheduler

Most images look very similar and are comparable in quality. Again, it often comes down to your specific use case so a good approach is to run multiple different schedulers and compare the results.

##
Models

Models are loaded from the ModelMixin.from_pretrained() method, which downloads and caches the latest version of the model weights and configurations. If the latest files are available in the local cache, from_pretrained() reuses files in the cache instead of re-downloading them.

Models can be loaded from a subfolder with the `subfolder` argument. For example, the model weights for stable-diffusion-v1-5/stable-diffusion-v1-5 are stored in the unet subfolder.

    from diffusers import UNet2DConditionModel

    unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", use_safetensors=True)

They can also be directly loaded from a repository.

    from diffusers import UNet2DModel

    unet = UNet2DModel.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)

To load and save model variants, specify the `variant` argument in ModelMixin.from_pretrained() and ModelMixin.save_pretrained().

    from diffusers import UNet2DConditionModel

    unet = UNet2DConditionModel.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet", variant="non_ema", use_safetensors=True
    )
    unet.save_pretrained("./local-unet", variant="non_ema")

Use the `torch_dtype` argument in from_pretrained() to specify the dtype to load a model in.

    from diffusers import AutoModel

    unet = AutoModel.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", torch_dtype=torch.float16
    )

You can also use the torch.Tensor.to method to convert to the specified dtype on the fly. It converts all weights unlike the `torch_dtype` argument that respects the `_keep_in_fp32_modules`. This is important for models whose layers must remain in fp32 for numerical stability and best generation quality (see example here).

< > Update on GitHub