Repository: huggingface/diffusers
File: src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py (Python)

Code before:
from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_dpmsolver_sde import BrownianTreeNoiseSampler
from .scheduling_utils import SchedulerMixin, SchedulerOutput


class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
    """
    Implements a variant of `DPMSolverMultistepScheduler` with cosine schedule, proposed by Nichol and Dhariwal (2021).
    This scheduler was used in Stable Audio Open [1].

    [1] Evans, Parker, et al. "Stable Audio Open" https://huggingface.co/papers/2407.14358

    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    methods the library implements for all schedulers such as loading and saving.

    Args:
        sigma_min (`float`, defaults to `0.3`):
            Minimum noise magnitude in the sigma schedule. This was set to 0.3 in Stable Audio Open [1].
        sigma_max (`float`, defaults to `500`):
            Maximum noise magnitude in the sigma schedule. This was set to 500 in Stable Audio Open [1].
        sigma_data (`float`, defaults to `1.0`):
            The standard deviation of the data distribution. This is set to 1.0 in Stable Audio Open [1].
        sigma_schedule (`str`, defaults to `"exponential"`):
            Sigma schedule to compute the `sigmas`. Must be one of `"exponential"` or `"karras"`. The exponential
            schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. The Karras schedule is
            introduced in the EDM paper (https://huggingface.co/papers/2206.00364).
        num_train_timesteps (`int`, defaults to `1000`):
            The number of diffusion steps to train the model.
        solver_order (`int`, defaults to `2`):
            The DPMSolver order which can be `1` or `2`. It is recommended to use `solver_order=2`.
        prediction_type (`str`, defaults to `"v_prediction"`):
            Prediction type of the scheduler function. Must be one of `"epsilon"` (predicts the noise of the diffusion
            process), `"sample"` (directly predicts the noisy sample), or `"v_prediction"` (see section 2.4 of [Imagen
            Video](https://huggingface.co/papers/2210.02303) paper).
        rho (`float`, defaults to `7.0`):
            The parameter for calculating the Karras sigma schedule from the EDM
            [paper](https://huggingface.co/papers/2206.00364).
        solver_type (`str`, defaults to `"midpoint"`):
            Solver type for the second-order solver. Must be one of `"midpoint"` or `"heun"`. The solver type slightly
            affects the sample quality, especially for a small number of steps. It is recommended to use `"midpoint"`.
        lower_order_final (`bool`, defaults to `True`):
            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
        euler_at_final (`bool`, defaults to `False`):
            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
            steps, but sometimes may result in blurring.
        final_sigmas_type (`str`, defaults to `"zero"`):
            The final `sigma` value for the noise schedule during the sampling process. Must be one of `"zero"` or
            `"sigma_min"`. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If
            `"zero"`, the final sigma is set to 0.

Review comment (suggestion):
```suggestion
            introduced in the [EDM](https://huggingface.co/papers/2206.00364) paper.
```

Code after:
from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_dpmsolver_sde import BrownianTreeNoiseSampler
from .scheduling_utils import SchedulerMixin, SchedulerOutput


class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
    """
    Implements a variant of `DPMSolverMultistepScheduler` with cosine schedule, proposed by Nichol and Dhariwal (2021).
    This scheduler was used in Stable Audio Open [1].

    [1] Evans, Parker, et al. "Stable Audio Open" https://huggingface.co/papers/2407.14358

    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    methods the library implements for all schedulers such as loading and saving.

    Args:
        sigma_min (`float`, defaults to `0.3`):
            Minimum noise magnitude in the sigma schedule. This was set to 0.3 in Stable Audio Open [1].
        sigma_max (`float`, defaults to `500`):
            Maximum noise magnitude in the sigma schedule. This was set to 500 in Stable Audio Open [1].
        sigma_data (`float`, defaults to `1.0`):
            The standard deviation of the data distribution. This is set to 1.0 in Stable Audio Open [1].
        sigma_schedule (`str`, defaults to `"exponential"`):
            Sigma schedule to compute the `sigmas`. Must be one of `"exponential"` or `"karras"`. The exponential
            schedule was incorporated in [stabilityai/cosxl](https://huggingface.co/stabilityai/cosxl). The Karras
            schedule is introduced in the [EDM](https://huggingface.co/papers/2206.00364) paper.
        num_train_timesteps (`int`, defaults to `1000`):
            The number of diffusion steps to train the model.
        solver_order (`int`, defaults to `2`):
            The DPMSolver order which can be `1` or `2`. It is recommended to use `solver_order=2`.
        prediction_type (`str`, defaults to `"v_prediction"`):
            Prediction type of the scheduler function. Must be one of `"epsilon"` (predicts the noise of the diffusion
            process), `"sample"` (directly predicts the noisy sample), or `"v_prediction"` (see section 2.4 of [Imagen
            Video](https://huggingface.co/papers/2210.02303) paper).
        rho (`float`, defaults to `7.0`):
            The parameter for calculating the Karras sigma schedule from the EDM
            [paper](https://huggingface.co/papers/2206.00364).
        solver_type (`str`, defaults to `"midpoint"`):
            Solver type for the second-order solver. Must be one of `"midpoint"` or `"heun"`. The solver type slightly
            affects the sample quality, especially for a small number of steps. It is recommended to use `"midpoint"`.
        lower_order_final (`bool`, defaults to `True`):
            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
        euler_at_final (`bool`, defaults to `False`):
            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
            steps, but sometimes may result in blurring.
        final_sigmas_type (`str`, defaults to `"zero"`):
            The final `sigma` value for the noise schedule during the sampling process. Must be one of `"zero"` or
            `"sigma_min"`. If `"sigma_min"`, the final sigma is the same as the last sigma in the training schedule. If
            `"zero"`, the final sigma is set to 0.