Schedulers

🤗 Diffusers provides many scheduler functions for the diffusion process. A scheduler takes a model's output (the sample which the diffusion process is iterating on) and a timestep to return a denoised sample. The timestep is important because it dictates where in the diffusion process the step is; data is generated by iterating forward n timesteps and inference occurs by propagating backward through the timesteps. Based on the timestep, a scheduler may be discrete in which case the timestep is an int or continuous in which case the timestep is a float.

Depending on the context, a scheduler defines how to iteratively add noise to an image or how to update a sample based on a model's output:

• during training, a scheduler adds noise (there are different algorithms for how to add noise) to a sample to train a diffusion model
• during inference, a scheduler defines how to update a sample based on a pretrained model's output

Many schedulers are implemented from the k-diffusion library by Katherine Crowson, and they're also widely used in A1111. To help you map the schedulers from k-diffusion and A1111 to the schedulers in 🤗 Diffusers, take a look at the table below:

A1111/k-diffusion	🤗 Diffusers	Usage
DPM++ 2M	DPMSolverMultistepScheduler
DPM++ 2M Karras	DPMSolverMultistepScheduler	init with use_karras_sigmas=True
DPM++ 2M SDE	DPMSolverMultistepScheduler	init with algorithm_type="sde-dpmsolver++"
DPM++ 2M SDE Karras	DPMSolverMultistepScheduler	init with use_karras_sigmas=True and algorithm_type="sde-dpmsolver++"
DPM++ 2S a	N/A	very similar to DPMSolverSinglestepScheduler
DPM++ 2S a Karras	N/A	very similar to DPMSolverSinglestepScheduler(use_karras_sigmas=True, ...)
DPM++ SDE	DPMSolverSinglestepScheduler
DPM++ SDE Karras	DPMSolverSinglestepScheduler	init with use_karras_sigmas=True
DPM2	KDPM2DiscreteScheduler
DPM2 Karras	KDPM2DiscreteScheduler	init with use_karras_sigmas=True
DPM2 a	KDPM2AncestralDiscreteScheduler
DPM2 a Karras	KDPM2AncestralDiscreteScheduler	init with use_karras_sigmas=True
DPM adaptive	N/A
DPM fast	N/A
Euler	EulerDiscreteScheduler
Euler a	EulerAncestralDiscreteScheduler
Heun	HeunDiscreteScheduler
LMS	LMSDiscreteScheduler
LMS Karras	LMSDiscreteScheduler	init with use_karras_sigmas=True
N/A	DEISMultistepScheduler
N/A	UniPCMultistepScheduler

All schedulers are built from the base SchedulerMixin class which implements low level utilities shared by all schedulers.

mindone.diffusers.SchedulerMixin

Bases: PushToHubMixin

Base class for all schedulers.

[SchedulerMixin] contains common functions shared by all schedulers such as general loading and saving functionalities.

[ConfigMixin] takes care of storing the configuration attributes (like num_train_timesteps) that are passed to the scheduler's __init__ function, and the attributes can be accessed by scheduler.config.num_train_timesteps.

Class attributes

• _compatibles (List[str]) -- A list of scheduler classes that are compatible with the parent scheduler class. Use [~ConfigMixin.from_config] to load a different compatible scheduler class (should be overridden by parent class).

Source code in mindone/diffusers/schedulers/scheduling_utils.py

@ms.jit_class
class SchedulerMixin(PushToHubMixin):
    """
    Base class for all schedulers.

    [`SchedulerMixin`] contains common functions shared by all schedulers such as general loading and saving
    functionalities.

    [`ConfigMixin`] takes care of storing the configuration attributes (like `num_train_timesteps`) that are passed to
    the scheduler's `__init__` function, and the attributes can be accessed by `scheduler.config.num_train_timesteps`.

    Class attributes:
        - **_compatibles** (`List[str]`) -- A list of scheduler classes that are compatible with the parent scheduler
          class. Use [`~ConfigMixin.from_config`] to load a different compatible scheduler class (should be overridden
          by parent class).
    """

    config_name = SCHEDULER_CONFIG_NAME
    _compatibles = []
    has_compatibles = True

    @classmethod
    @validate_hf_hub_args
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
        subfolder: Optional[str] = None,
        return_unused_kwargs=False,
        **kwargs,
    ):
        r"""
        Instantiate a scheduler from a pre-defined JSON configuration file in a local directory or Hub repository.

        Parameters:
            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
                Can be either:

                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
                      the Hub.
                    - A path to a *directory* (for example `./my_model_directory`) containing the scheduler
                      configuration saved with [`~SchedulerMixin.save_pretrained`].
            subfolder (`str`, *optional*):
                The subfolder location of a model file within a larger model repository on the Hub or locally.
            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
                Whether kwargs that are not consumed by the Python class should be returned or not.
            cache_dir (`Union[str, os.PathLike]`, *optional*):
                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
                is not used.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.

            proxies (`Dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            output_loading_info(`bool`, *optional*, defaults to `False`):
                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
            local_files_only(`bool`, *optional*, defaults to `False`):
                Whether to only load local model weights and configuration files or not. If set to `True`, the model
                won't be downloaded from the Hub.
            token (`str` or *bool*, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
                `diffusers-cli login` (stored in `~/.huggingface`) is used.
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
                allowed by Git.

        <Tip>

        To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with
        `huggingface-cli login`. You can also activate the special
        ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a
        firewalled environment.

        </Tip>

        """
        config, kwargs, commit_hash = cls.load_config(
            pretrained_model_name_or_path=pretrained_model_name_or_path,
            subfolder=subfolder,
            return_unused_kwargs=True,
            return_commit_hash=True,
            **kwargs,
        )
        return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)

    def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
        """
        Save a scheduler configuration object to a directory so that it can be reloaded using the
        [`~SchedulerMixin.from_pretrained`] class method.

        Args:
            save_directory (`str` or `os.PathLike`):
                Directory where the configuration JSON file will be saved (will be created if it does not exist).
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            kwargs (`Dict[str, Any]`, *optional*):
                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)

    @property
    def compatibles(self):
        """
        Returns all schedulers that are compatible with this scheduler

        Returns:
            `List[SchedulerMixin]`: List of compatible schedulers
        """
        return self._get_compatibles()

    @classmethod
    def _get_compatibles(cls):
        compatible_classes_str = list(set([cls.__name__] + cls._compatibles))
        diffusers_library = maybe_import_module_in_mindone(__name__.split(".")[1])
        compatible_classes = [
            getattr(diffusers_library, c) for c in compatible_classes_str if hasattr(diffusers_library, c)
        ]
        return compatible_classes

mindone.diffusers.SchedulerMixin.from_pretrained(pretrained_model_name_or_path=None, subfolder=None, return_unused_kwargs=False, **kwargs) classmethod

Instantiate a scheduler from a pre-defined JSON configuration file in a local directory or Hub repository.

PARAMETER	DESCRIPTION
pretrained_model_name_or_path	Can be either: - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on the Hub. - A path to a *directory* (for example `./my_model_directory`) containing the scheduler configuration saved with [`~SchedulerMixin.save_pretrained`]. TYPE: `str` or `os.PathLike`, *optional* DEFAULT: None
subfolder	The subfolder location of a model file within a larger model repository on the Hub or locally. TYPE: `str`, *optional* DEFAULT: None
return_unused_kwargs	Whether kwargs that are not consumed by the Python class should be returned or not. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
cache_dir	Path to a directory where a downloaded pretrained model configuration is cached if the standard cache is not used. TYPE: `Union[str, os.PathLike]`, *optional*
force_download	Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist. TYPE: `bool`, *optional*, defaults to `False`
proxies	A dictionary of proxy servers to use by protocol or endpoint, for example, {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}. The proxies are used on each request. TYPE: `Dict[str, str]`, *optional*
output_loading_info(`bool`,	Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. TYPE: *optional*, defaults to `False`
local_files_only(`bool`,	Whether to only load local model weights and configuration files or not. If set to True, the model won't be downloaded from the Hub. TYPE: *optional*, defaults to `False`
token	The token to use as HTTP bearer authorization for remote files. If True, the token generated from diffusers-cli login (stored in ~/.huggingface) is used. TYPE: `str` or *bool*, *optional*
revision	The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier allowed by Git. TYPE: `str`, *optional*, defaults to `"main"`

To use private or gated models, log-in with huggingface-cli login. You can also activate the special "offline-mode" to use this method in a firewalled environment.

Source code in mindone/diffusers/schedulers/scheduling_utils.py

@classmethod
@validate_hf_hub_args
def from_pretrained(
    cls,
    pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
    subfolder: Optional[str] = None,
    return_unused_kwargs=False,
    **kwargs,
):
    r"""
    Instantiate a scheduler from a pre-defined JSON configuration file in a local directory or Hub repository.

    Parameters:
        pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
            Can be either:

                - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
                  the Hub.
                - A path to a *directory* (for example `./my_model_directory`) containing the scheduler
                  configuration saved with [`~SchedulerMixin.save_pretrained`].
        subfolder (`str`, *optional*):
            The subfolder location of a model file within a larger model repository on the Hub or locally.
        return_unused_kwargs (`bool`, *optional*, defaults to `False`):
            Whether kwargs that are not consumed by the Python class should be returned or not.
        cache_dir (`Union[str, os.PathLike]`, *optional*):
            Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
            is not used.
        force_download (`bool`, *optional*, defaults to `False`):
            Whether or not to force the (re-)download of the model weights and configuration files, overriding the
            cached versions if they exist.

        proxies (`Dict[str, str]`, *optional*):
            A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
            'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
        output_loading_info(`bool`, *optional*, defaults to `False`):
            Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
        local_files_only(`bool`, *optional*, defaults to `False`):
            Whether to only load local model weights and configuration files or not. If set to `True`, the model
            won't be downloaded from the Hub.
        token (`str` or *bool*, *optional*):
            The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
            `diffusers-cli login` (stored in `~/.huggingface`) is used.
        revision (`str`, *optional*, defaults to `"main"`):
            The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
            allowed by Git.

    <Tip>

    To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with
    `huggingface-cli login`. You can also activate the special
    ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a
    firewalled environment.

    </Tip>

    """
    config, kwargs, commit_hash = cls.load_config(
        pretrained_model_name_or_path=pretrained_model_name_or_path,
        subfolder=subfolder,
        return_unused_kwargs=True,
        return_commit_hash=True,
        **kwargs,
    )
    return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)

mindone.diffusers.SchedulerMixin.save_pretrained(save_directory, push_to_hub=False, **kwargs)

Save a scheduler configuration object to a directory so that it can be reloaded using the [~SchedulerMixin.from_pretrained] class method.

PARAMETER	DESCRIPTION
save_directory	Directory where the configuration JSON file will be saved (will be created if it does not exist). TYPE: `str` or `os.PathLike`
push_to_hub	Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the repository you want to push to with repo_id (will default to the name of save_directory in your namespace). TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
kwargs	Additional keyword arguments passed along to the [~utils.PushToHubMixin.push_to_hub] method. TYPE: `Dict[str, Any]`, *optional* DEFAULT: {}

Source code in mindone/diffusers/schedulers/scheduling_utils.py

def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
    """
    Save a scheduler configuration object to a directory so that it can be reloaded using the
    [`~SchedulerMixin.from_pretrained`] class method.

    Args:
        save_directory (`str` or `os.PathLike`):
            Directory where the configuration JSON file will be saved (will be created if it does not exist).
        push_to_hub (`bool`, *optional*, defaults to `False`):
            Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
            repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
            namespace).
        kwargs (`Dict[str, Any]`, *optional*):
            Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
    """
    self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)

mindone.diffusers.schedulers.scheduling_utils.SchedulerOutput dataclass

Bases: BaseOutput

Base class for the output of a scheduler's step function.

PARAMETER	DESCRIPTION
prev_sample	Computed sample (x_{t-1}) of previous timestep. prev_sample should be used as next model input in the denoising loop. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images

Source code in mindone/diffusers/schedulers/scheduling_utils.py

@dataclass
class SchedulerOutput(BaseOutput):
    """
    Base class for the output of a scheduler's `step` function.

    Args:
        prev_sample (`ms.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
    """

    prev_sample: ms.Tensor

KarrasDiffusionSchedulers

[KarrasDiffusionSchedulers] are a broad generalization of schedulers in 🤗 Diffusers. The schedulers in this class are distinguished at a high level by their noise sampling strategy, the type of network and scaling, the training strategy, and how the loss is weighed.

The different schedulers in this class, depending on the ordinary differential equations (ODE) solver type, fall into the above taxonomy and provide a good abstraction for the design of the main schedulers implemented in 🤗 Diffusers. The schedulers in this class are given here.

mindone.diffusers.utils.PushToHubMixin

A Mixin to push a model, scheduler, or pipeline to the Hugging Face Hub.

Source code in mindone/diffusers/utils/hub_utils.py

class PushToHubMixin:
    """
    A Mixin to push a model, scheduler, or pipeline to the Hugging Face Hub.
    """

    def _upload_folder(
        self,
        working_dir: Union[str, os.PathLike],
        repo_id: str,
        token: Optional[str] = None,
        commit_message: Optional[str] = None,
        create_pr: bool = False,
    ):
        """
        Uploads all files in `working_dir` to `repo_id`.
        """
        if commit_message is None:
            if "Model" in self.__class__.__name__:
                commit_message = "Upload model"
            elif "Scheduler" in self.__class__.__name__:
                commit_message = "Upload scheduler"
            else:
                commit_message = f"Upload {self.__class__.__name__}"

        logger.info(f"Uploading the files of {working_dir} to {repo_id}.")
        return upload_folder(
            repo_id=repo_id, folder_path=working_dir, token=token, commit_message=commit_message, create_pr=create_pr
        )

    def push_to_hub(
        self,
        repo_id: str,
        commit_message: Optional[str] = None,
        private: Optional[bool] = None,
        token: Optional[str] = None,
        create_pr: bool = False,
        safe_serialization: bool = True,
        variant: Optional[str] = None,
    ) -> str:
        """
        Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub.

        Parameters:
            repo_id (`str`):
                The name of the repository you want to push your model, scheduler, or pipeline files to. It should
                contain your organization name when pushing to an organization. `repo_id` can also be a path to a local
                directory.
            commit_message (`str`, *optional*):
                Message to commit while pushing. Default to `"Upload {object}"`.
            private (`bool`, *optional*):
                Whether or not the repository created should be private.
            token (`str`, *optional*):
                The token to use as HTTP bearer authorization for remote files. The token generated when running
                `huggingface-cli login` (stored in `~/.huggingface`).
            create_pr (`bool`, *optional*, defaults to `False`):
                Whether or not to create a PR with the uploaded files or directly commit.
            safe_serialization (`bool`, *optional*, defaults to `True`):
                Whether or not to convert the model weights to the `safetensors` format.
            variant (`str`, *optional*):
                If specified, weights are saved in the format `pytorch_model.<variant>.bin`.

        Examples:

        ```python
        from mindone.diffusers import UNet2DConditionModel

        unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet")

        # Push the `unet` to your namespace with the name "my-finetuned-unet".
        unet.push_to_hub("my-finetuned-unet")

        # Push the `unet` to an organization with the name "my-finetuned-unet".
        unet.push_to_hub("your-org/my-finetuned-unet")
        ```
        """
        repo_id = create_repo(repo_id, private=private, token=token, exist_ok=True).repo_id

        # Create a new empty model card and eventually tag it
        model_card = load_or_create_model_card(repo_id, token=token)
        model_card = populate_model_card(model_card)

        # Save all files.
        save_kwargs = {"safe_serialization": safe_serialization}
        if "Scheduler" not in self.__class__.__name__:
            save_kwargs.update({"variant": variant})

        with tempfile.TemporaryDirectory() as tmpdir:
            self.save_pretrained(tmpdir, **save_kwargs)

            # Update model card if needed:
            model_card.save(os.path.join(tmpdir, "README.md"))

            return self._upload_folder(
                tmpdir,
                repo_id,
                token=token,
                commit_message=commit_message,
                create_pr=create_pr,
            )

mindone.diffusers.utils.PushToHubMixin.push_to_hub(repo_id, commit_message=None, private=None, token=None, create_pr=False, safe_serialization=True, variant=None)

Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub.

PARAMETER	DESCRIPTION
repo_id	The name of the repository you want to push your model, scheduler, or pipeline files to. It should contain your organization name when pushing to an organization. repo_id can also be a path to a local directory. TYPE: `str`
commit_message	Message to commit while pushing. Default to "Upload {object}". TYPE: `str`, *optional* DEFAULT: None
private	Whether or not the repository created should be private. TYPE: `bool`, *optional* DEFAULT: None
token	The token to use as HTTP bearer authorization for remote files. The token generated when running huggingface-cli login (stored in ~/.huggingface). TYPE: `str`, *optional* DEFAULT: None
create_pr	Whether or not to create a PR with the uploaded files or directly commit. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
safe_serialization	Whether or not to convert the model weights to the safetensors format. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
variant	If specified, weights are saved in the format pytorch_model.<variant>.bin. TYPE: `str`, *optional* DEFAULT: None

from mindone.diffusers import UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet")

# Push the `unet` to your namespace with the name "my-finetuned-unet".
unet.push_to_hub("my-finetuned-unet")

# Push the `unet` to an organization with the name "my-finetuned-unet".
unet.push_to_hub("your-org/my-finetuned-unet")

Source code in mindone/diffusers/utils/hub_utils.py

def push_to_hub(
    self,
    repo_id: str,
    commit_message: Optional[str] = None,
    private: Optional[bool] = None,
    token: Optional[str] = None,
    create_pr: bool = False,
    safe_serialization: bool = True,
    variant: Optional[str] = None,
) -> str:
    """
    Upload model, scheduler, or pipeline files to the 🤗 Hugging Face Hub.

    Parameters:
        repo_id (`str`):
            The name of the repository you want to push your model, scheduler, or pipeline files to. It should
            contain your organization name when pushing to an organization. `repo_id` can also be a path to a local
            directory.
        commit_message (`str`, *optional*):
            Message to commit while pushing. Default to `"Upload {object}"`.
        private (`bool`, *optional*):
            Whether or not the repository created should be private.
        token (`str`, *optional*):
            The token to use as HTTP bearer authorization for remote files. The token generated when running
            `huggingface-cli login` (stored in `~/.huggingface`).
        create_pr (`bool`, *optional*, defaults to `False`):
            Whether or not to create a PR with the uploaded files or directly commit.
        safe_serialization (`bool`, *optional*, defaults to `True`):
            Whether or not to convert the model weights to the `safetensors` format.
        variant (`str`, *optional*):
            If specified, weights are saved in the format `pytorch_model.<variant>.bin`.

    Examples:

    ```python
    from mindone.diffusers import UNet2DConditionModel

    unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet")

    # Push the `unet` to your namespace with the name "my-finetuned-unet".
    unet.push_to_hub("my-finetuned-unet")

    # Push the `unet` to an organization with the name "my-finetuned-unet".
    unet.push_to_hub("your-org/my-finetuned-unet")
    ```
    """
    repo_id = create_repo(repo_id, private=private, token=token, exist_ok=True).repo_id

    # Create a new empty model card and eventually tag it
    model_card = load_or_create_model_card(repo_id, token=token)
    model_card = populate_model_card(model_card)

    # Save all files.
    save_kwargs = {"safe_serialization": safe_serialization}
    if "Scheduler" not in self.__class__.__name__:
        save_kwargs.update({"variant": variant})

    with tempfile.TemporaryDirectory() as tmpdir:
        self.save_pretrained(tmpdir, **save_kwargs)

        # Update model card if needed:
        model_card.save(os.path.join(tmpdir, "README.md"))

        return self._upload_folder(
            tmpdir,
            repo_id,
            token=token,
            commit_message=commit_message,
            create_pr=create_pr,
        )