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 mindspore as ms
from mindone.diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", mindspore_dtype=ms.float16, use_safetensors=True
)
You can see what scheduler this pipeline uses with the pipeline.scheduler attribute.
pipeline.scheduler
PNDMScheduler {
"_class_name": "PNDMScheduler",
"_diffusers_version": "0.29.2",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": false,
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"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 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 mindone.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, mindspore_dtype=ms.float16, use_safetensors=True
)
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 mindspore as ms
from mindone.diffusers import DiffusionPipeline
import numpy as np
pipeline = DiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", mindspore_dtype=ms.float16, use_safetensors=True
)
prompt = "A photograph of an astronaut riding a horse on Mars, high resolution, high definition."
generator = np.random.Generator(np.random.PCG64(8))
To change the pipelines scheduler, use the from_config method to load a different scheduler's pipeline.scheduler.config into the pipeline.
LMSDiscreteScheduler typically generates higher quality images than the default scheduler.
from mindone.diffusers import LMSDiscreteScheduler
pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator)[0][0]
image
EulerDiscreteScheduler can generate higher quality images in just 30 steps.
from mindone.diffusers import EulerDiscreteScheduler
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator)[0][0]
image
EulerAncestralDiscreteScheduler can generate higher quality images in just 30 steps.
from mindone.diffusers import EulerAncestralDiscreteScheduler
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator)[0][0]
image
DPMSolverMultistepScheduler provides a balance between speed and quality and can generate higher quality images in just 20 steps.
from mindone.diffusers import DPMSolverMultistepScheduler
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator)[0][0]
image
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 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 mindone.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 mindone.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 mindone.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")