From 41e3bf2e9a1e0037eac2dcd400dfbe6fc7e89ff0 Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Thu, 27 Oct 2022 12:42:04 +0100 Subject: [PATCH 01/27] k-diffusion-euler --- src/diffusers/__init__.py | 2 + .../pipeline_stable_diffusion.py | 6 +- .../pipeline_stable_diffusion_img2img.py | 8 +- src/diffusers/schedulers/__init__.py | 2 + .../scheduling_euler_ancestral_discrete.py | 285 +++++++++++++++++ .../schedulers/scheduling_euler_discrete.py | 293 ++++++++++++++++++ .../stable_diffusion/test_stable_diffusion.py | 92 ++++++ tests/test_scheduler.py | 113 +++++++ 8 files changed, 797 insertions(+), 4 deletions(-) create mode 100644 src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py create mode 100644 src/diffusers/schedulers/scheduling_euler_discrete.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2c531cf8cee0..80e1edc4ab4f 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -53,6 +53,8 @@ if is_torch_available() and is_scipy_available(): from .schedulers import LMSDiscreteScheduler + from .schedulers import EulerAncestralDiscreteScheduler + from .schedulers import EulerDiscreteScheduler else: from .utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index cf4c5c5fdeca..7eac14c39732 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -9,7 +9,8 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler +from ...schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, \ +EulerDiscreteScheduler, EulerAncestralDiscreteScheduler from ...utils import deprecate, logging from . import StableDiffusionPipelineOutput from .safety_checker import StableDiffusionSafetyChecker @@ -52,7 +53,8 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, + EulerDiscreteScheduler, EulerAncestralDiscreteScheduler], safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPFeatureExtractor, ): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index db8134834564..dbf3ed738cf4 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -10,7 +10,8 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler +from ...schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, \ +EulerDiscreteScheduler, EulerAncestralDiscreteScheduler from ...utils import deprecate, logging from . import StableDiffusionPipelineOutput from .safety_checker import StableDiffusionSafetyChecker @@ -63,7 +64,10 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], + scheduler: Union[ + DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, + EulerDiscreteScheduler, EulerAncestralDiscreteScheduler + ], safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPFeatureExtractor, ): diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index f179acde09e2..6db70be54bb6 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -42,5 +42,7 @@ if is_scipy_available() and is_torch_available(): from .scheduling_lms_discrete import LMSDiscreteScheduler + from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler + from .scheduling_euler_discrete import EulerDiscreteScheduler else: from ..utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py new file mode 100644 index 000000000000..b3b2866bc715 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -0,0 +1,285 @@ +# Copyright 2022 Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import numpy as np +import torch + +from scipy import integrate + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput, deprecate +from .scheduling_utils import SchedulerMixin + + +@dataclass +# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete +class EulerAncestralDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's step function output. + + Args: + prev_sample (`torch.FloatTensor` 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. + pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + The predicted denoised sample (x_{0}) based on the model output from the current timestep. + `pred_original_sample` can be used to preview progress or for guidance. + """ + + prev_sample: torch.FloatTensor + pred_original_sample: Optional[torch.FloatTensor] = None + + +class EulerAncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by + Katherine Crowson: + https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 + + [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` + function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. + [`~ConfigMixin`] also provides general loading and saving functionality via the [`~ConfigMixin.save_config`] and + [`~ConfigMixin.from_config`] functions. + + Args: + num_train_timesteps (`int`): number of diffusion steps used to train the model. + beta_start (`float`): the starting `beta` value of inference. + beta_end (`float`): the final `beta` value. + beta_schedule (`str`): + the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + `linear` or `scaled_linear`. + trained_betas (`np.ndarray`, optional): + option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. + + """ + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + beta_schedule: str = "linear", + trained_betas: Optional[np.ndarray] = None, + ): + if trained_betas is not None: + self.betas = torch.from_numpy(trained_betas) + elif beta_schedule == "linear": + self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = ( + torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + ) + else: + raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas) + + # standard deviation of the initial noise distribution + self.init_noise_sigma = self.sigmas.max() + + # setable values + self.num_inference_steps = None + timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() + self.timesteps = torch.from_numpy(timesteps) + self.derivatives = [] + self.is_scale_input_called = False + + def scale_model_input( + self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] + ) -> torch.FloatTensor: + """ + Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. + + Args: + sample (`torch.FloatTensor`): input sample + timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain + + Returns: + `torch.FloatTensor`: scaled input sample + """ + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + step_index = (self.timesteps == timestep).nonzero().item() + sigma = self.sigmas[step_index] + sample = sample / ((sigma**2 + 1) ** 0.5) + self.is_scale_input_called = True + return sample + + def get_lms_coefficient(self, order, t, current_order): + """ + Compute a linear multistep coefficient. + + Args: + order (TODO): + t (TODO): + current_order (TODO): + """ + + def lms_derivative(tau): + prod = 1.0 + for k in range(order): + if current_order == k: + continue + prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) + return prod + + integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] + + return integrated_coeff + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + """ + Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. + + Args: + num_inference_steps (`int`): + the number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, optional): + the device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + self.num_inference_steps = num_inference_steps + + timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy() + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device) + + self.derivatives = [] + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + order: int = 4, + return_dict: bool = True, + ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): direct output from learned diffusion model. + timestep (`float`): current timestep in the diffusion chain. + sample (`torch.FloatTensor`): + current instance of sample being created by diffusion process. + order: coefficient for multi-step inference. + return_dict (`bool`): option for returning tuple rather than EulerAncestralDiscreteSchedulerOutput class + + Returns: + [`~schedulers.scheduling_utils.EulerAncestralDiscreteSchedulerOutput`] or `tuple`: + [`~schedulers.scheduling_utils.EulerAncestralDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. + When returning a tuple, the first element is the sample tensor. + + """ + if not self.is_scale_input_called: + warnings.warn( + "The `scale_model_input` function should be called before `step` to ensure correct denoising. " + "See `StableDiffusionPipeline` for a usage example." + ) + + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + deprecate( + "timestep as an index", + "0.8.0", + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerAncestralDiscreteScheduler.step()` will not be supported in future versions. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep.", + standard_warn=False, + ) + step_index = timestep + else: + step_index = (self.timesteps == timestep).nonzero().item() + sigma = self.sigmas[step_index] + + # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise + pred_original_sample = sample - sigma * model_output + sigma_from = self.sigmas[step_index] + sigma_to = self.sigmas[step_index + 1] + sigma_up = (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5 + sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5 + # 2. Convert to an ODE derivative + derivative = (sample - pred_original_sample) / sigma + self.derivatives.append(derivative) + if len(self.derivatives) > order: + self.derivatives.pop(0) + dt = sigma_down - sigma + + prev_sample = sample + derivative * dt + + prev_sample = prev_sample + torch.randn_like(prev_sample) * sigma_up + + if not return_dict: + return (prev_sample,) + + return EulerAncestralDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) + + def add_noise( + self, + original_samples: torch.FloatTensor, + noise: torch.FloatTensor, + timesteps: torch.FloatTensor, + ) -> torch.FloatTensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + self.sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + # mps does not support float64 + self.timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) + timesteps = timesteps.to(original_samples.device, dtype=torch.float32) + else: + self.timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + schedule_timesteps = self.timesteps + + if isinstance(timesteps, torch.IntTensor) or isinstance(timesteps, torch.LongTensor): + deprecate( + "timesteps as indices", + "0.8.0", + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerAncestralDiscreteScheduler.add_noise()` will not be supported in future versions. Make sure to" + " pass values from `scheduler.timesteps` as timesteps.", + standard_warn=False, + ) + step_indices = timesteps + else: + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = self.sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + noisy_samples = original_samples + noise * sigma + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py new file mode 100644 index 000000000000..a76c31bbd392 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -0,0 +1,293 @@ +# Copyright 2022 Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import numpy as np +import torch + +from scipy import integrate + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput, deprecate +from .scheduling_utils import SchedulerMixin + + +@dataclass +# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete +class EulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's step function output. + + Args: + prev_sample (`torch.FloatTensor` 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. + pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + The predicted denoised sample (x_{0}) based on the model output from the current timestep. + `pred_original_sample` can be used to preview progress or for guidance. + """ + + prev_sample: torch.FloatTensor + pred_original_sample: Optional[torch.FloatTensor] = None + + +class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by + Katherine Crowson: + https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 + + [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` + function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. + [`~ConfigMixin`] also provides general loading and saving functionality via the [`~ConfigMixin.save_config`] and + [`~ConfigMixin.from_config`] functions. + + Args: + num_train_timesteps (`int`): number of diffusion steps used to train the model. + beta_start (`float`): the starting `beta` value of inference. + beta_end (`float`): the final `beta` value. + beta_schedule (`str`): + the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + `linear` or `scaled_linear`. + trained_betas (`np.ndarray`, optional): + option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc. + + """ + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + beta_schedule: str = "linear", + trained_betas: Optional[np.ndarray] = None, + ): + if trained_betas is not None: + self.betas = torch.from_numpy(trained_betas) + elif beta_schedule == "linear": + self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = ( + torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + ) + else: + raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas) + + # standard deviation of the initial noise distribution + self.init_noise_sigma = self.sigmas.max() + + # setable values + self.num_inference_steps = None + timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() + self.timesteps = torch.from_numpy(timesteps) + self.derivatives = [] + self.is_scale_input_called = False + + def scale_model_input( + self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] + ) -> torch.FloatTensor: + """ + Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. + + Args: + sample (`torch.FloatTensor`): input sample + timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain + + Returns: + `torch.FloatTensor`: scaled input sample + """ + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + step_index = (self.timesteps == timestep).nonzero().item() + sigma = self.sigmas[step_index] + sample = sample / ((sigma**2 + 1) ** 0.5) + self.is_scale_input_called = True + return sample + + def get_lms_coefficient(self, order, t, current_order): + """ + Compute a linear multistep coefficient. + + Args: + order (TODO): + t (TODO): + current_order (TODO): + """ + + def lms_derivative(tau): + prod = 1.0 + for k in range(order): + if current_order == k: + continue + prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) + return prod + + integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] + + return integrated_coeff + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + """ + Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. + + Args: + num_inference_steps (`int`): + the number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, optional): + the device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + self.num_inference_steps = num_inference_steps + + timesteps = np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps, dtype=float)[::-1].copy() + sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) + sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas) + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device) + + self.derivatives = [] + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + order: int = 4, + s_churn: float = 0., + s_tmin: float = 0., + s_tmax: float = float('inf'), + s_noise: float = 1., + return_dict: bool = True, + ) -> Union[EulerDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): direct output from learned diffusion model. + timestep (`float`): current timestep in the diffusion chain. + sample (`torch.FloatTensor`): + current instance of sample being created by diffusion process. + order: coefficient for multi-step inference. + s_churn (`float`) + s_tmin (`float`) + s_tmax (`float`) + s_noise (`float`) + return_dict (`bool`): option for returning tuple rather than EulerDiscreteSchedulerOutput class + + Returns: + [`~schedulers.scheduling_utils.EulerDiscreteSchedulerOutput`] or `tuple`: + [`~schedulers.scheduling_utils.EulerDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. + When returning a tuple, the first element is the sample tensor. + + """ + if not self.is_scale_input_called: + warnings.warn( + "The `scale_model_input` function should be called before `step` to ensure correct denoising. " + "See `StableDiffusionPipeline` for a usage example." + ) + + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + deprecate( + "timestep as an index", + "0.8.0", + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` will not be supported in future versions. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep.", + standard_warn=False, + ) + step_index = timestep + else: + step_index = (self.timesteps == timestep).nonzero().item() + sigma = self.sigmas[step_index] + gamma = min(s_churn / (len(self.sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigma <= s_tmax else 0. + eps = torch.randn_like(sample) * s_noise + sigma_hat = sigma * (gamma + 1) + if gamma > 0: + sample = sample + eps * (sigma_hat ** 2 - sigma ** 2) ** 0.5 + # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise + pred_original_sample = sample - sigma_hat * model_output + + # 2. Convert to an ODE derivative + derivative = (sample - pred_original_sample) / sigma_hat + self.derivatives.append(derivative) + if len(self.derivatives) > order: + self.derivatives.pop(0) + + dt = self.sigmas[step_index + 1] - sigma_hat + + prev_sample = sample + derivative * dt + + if not return_dict: + return (prev_sample,) + + return EulerDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) + + def add_noise( + self, + original_samples: torch.FloatTensor, + noise: torch.FloatTensor, + timesteps: torch.FloatTensor, + ) -> torch.FloatTensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + self.sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + # mps does not support float64 + self.timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) + timesteps = timesteps.to(original_samples.device, dtype=torch.float32) + else: + self.timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + schedule_timesteps = self.timesteps + + if isinstance(timesteps, torch.IntTensor) or isinstance(timesteps, torch.LongTensor): + deprecate( + "timesteps as indices", + "0.8.0", + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.add_noise()` will not be supported in future versions. Make sure to" + " pass values from `scheduler.timesteps` as timesteps.", + standard_warn=False, + ) + step_indices = timesteps + else: + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + + sigma = self.sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + noisy_samples = original_samples + noise * sigma + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index d7e6c362d1d7..610efcd1deb6 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -24,6 +24,8 @@ AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel, @@ -360,6 +362,96 @@ def test_stable_diffusion_k_lms(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + def test_stable_diffusion_k_euler_ancestral(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + unet = self.dummy_cond_unet + scheduler = EulerAncestralDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear") + vae = self.dummy_vae + bert = self.dummy_text_encoder + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + # make sure here that pndm scheduler skips prk + sd_pipe = StableDiffusionPipeline( + unet=unet, + scheduler=scheduler, + vae=vae, + text_encoder=bert, + tokenizer=tokenizer, + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) + + prompt = "A painting of a squirrel eating a burger" + generator = torch.Generator(device=device).manual_seed(0) + output = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") + + image = output.images + + generator = torch.Generator(device=device).manual_seed(0) + image_from_tuple = sd_pipe( + [prompt], + generator=generator, + guidance_scale=6.0, + num_inference_steps=2, + output_type="np", + return_dict=False, + )[0] + + image_slice = image[0, -3:, -3:, -1] + image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] + + assert image.shape == (1, 128, 128, 3) + expected_slice = np.array([0.5067, 0.4689, 0.4614, 0.5233, 0.4903, 0.5112, 0.524, 0.5069, 0.4785]) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + + def test_stable_diffusion_k_euler(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + unet = self.dummy_cond_unet + scheduler = EulerDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear") + vae = self.dummy_vae + bert = self.dummy_text_encoder + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + # make sure here that pndm scheduler skips prk + sd_pipe = StableDiffusionPipeline( + unet=unet, + scheduler=scheduler, + vae=vae, + text_encoder=bert, + tokenizer=tokenizer, + safety_checker=None, + feature_extractor=self.dummy_extractor, + ) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) + + prompt = "A painting of a squirrel eating a burger" + generator = torch.Generator(device=device).manual_seed(0) + output = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") + + image = output.images + + generator = torch.Generator(device=device).manual_seed(0) + image_from_tuple = sd_pipe( + [prompt], + generator=generator, + guidance_scale=6.0, + num_inference_steps=2, + output_type="np", + return_dict=False, + )[0] + + image_slice = image[0, -3:, -3:, -1] + image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] + + assert image.shape == (1, 128, 128, 3) + expected_slice = np.array([0.5067, 0.4689, 0.4614, 0.5233, 0.4903, 0.5112, 0.524, 0.5069, 0.4785]) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + def test_stable_diffusion_attention_chunk(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator unet = self.dummy_cond_unet diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 194fb66f663f..dbfc2290ba45 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -24,6 +24,8 @@ DDPMScheduler, IPNDMScheduler, LMSDiscreteScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, PNDMScheduler, ScoreSdeVeScheduler, ) @@ -932,6 +934,117 @@ def test_full_loop_no_noise(self): assert abs(result_sum.item() - 1006.388) < 1e-2 assert abs(result_mean.item() - 1.31) < 1e-3 +class EulerDiscreteSchedulerTest(SchedulerCommonTest): + scheduler_classes = (EulerDiscreteScheduler,) + num_inference_steps = 10 + + def get_scheduler_config(self, **kwargs): + config = { + "num_train_timesteps": 1100, + "beta_start": 0.0001, + "beta_end": 0.02, + "beta_schedule": "linear", + "trained_betas": None, + } + + config.update(**kwargs) + return config + + def test_timesteps(self): + for timesteps in [10, 50, 100, 1000]: + self.check_over_configs(num_train_timesteps=timesteps) + + def test_betas(self): + for beta_start, beta_end in zip([0.00001, 0.0001, 0.001], [0.0002, 0.002, 0.02]): + self.check_over_configs(beta_start=beta_start, beta_end=beta_end) + + def test_schedules(self): + for schedule in ["linear", "scaled_linear"]: + self.check_over_configs(beta_schedule=schedule) + + def test_time_indices(self): + for t in [0, 500, 800]: + self.check_over_forward(time_step=t) + + def test_full_loop_no_noise(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + scheduler.set_timesteps(self.num_inference_steps) + + model = self.dummy_model() + sample = self.dummy_sample_deter * scheduler.init_noise_sigma + + for i, t in enumerate(scheduler.timesteps): + sample = scheduler.scale_model_input(sample, t) + + model_output = model(sample, t) + + output = scheduler.step(model_output, t, sample) + sample = output.prev_sample + + result_sum = torch.sum(torch.abs(sample)) + result_mean = torch.mean(torch.abs(sample)) + + assert abs(result_sum.item() - 1006.388) < 1e-2 + assert abs(result_mean.item() - 1.31) < 1e-3 + +class EulerAncestralDiscreteSchedulerTest(SchedulerCommonTest): + scheduler_classes = (EulerAncestralDiscreteScheduler,) + num_inference_steps = 10 + + def get_scheduler_config(self, **kwargs): + config = { + "num_train_timesteps": 1100, + "beta_start": 0.0001, + "beta_end": 0.02, + "beta_schedule": "linear", + "trained_betas": None, + } + + config.update(**kwargs) + return config + + def test_timesteps(self): + for timesteps in [10, 50, 100, 1000]: + self.check_over_configs(num_train_timesteps=timesteps) + + def test_betas(self): + for beta_start, beta_end in zip([0.00001, 0.0001, 0.001], [0.0002, 0.002, 0.02]): + self.check_over_configs(beta_start=beta_start, beta_end=beta_end) + + def test_schedules(self): + for schedule in ["linear", "scaled_linear"]: + self.check_over_configs(beta_schedule=schedule) + + def test_time_indices(self): + for t in [0, 500, 800]: + self.check_over_forward(time_step=t) + + def test_full_loop_no_noise(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + scheduler.set_timesteps(self.num_inference_steps) + + model = self.dummy_model() + sample = self.dummy_sample_deter * scheduler.init_noise_sigma + + for i, t in enumerate(scheduler.timesteps): + sample = scheduler.scale_model_input(sample, t) + + model_output = model(sample, t) + + output = scheduler.step(model_output, t, sample) + sample = output.prev_sample + + result_sum = torch.sum(torch.abs(sample)) + result_mean = torch.mean(torch.abs(sample)) + + assert abs(result_sum.item() - 1006.388) < 1e-2 + assert abs(result_mean.item() - 1.31) < 1e-3 class IPNDMSchedulerTest(SchedulerCommonTest): scheduler_classes = (IPNDMScheduler,) From 9032c0b466fbe6140b3cb3d20649bbc7db50c42a Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Thu, 27 Oct 2022 12:49:01 +0100 Subject: [PATCH 02/27] make style make quality --- src/diffusers/__init__.py | 4 +--- .../pipeline_stable_diffusion.py | 14 ++++++++++---- .../pipeline_stable_diffusion_img2img.py | 12 ++++++++---- src/diffusers/schedulers/__init__.py | 2 +- .../scheduling_euler_ancestral_discrete.py | 12 +++++++----- .../schedulers/scheduling_euler_discrete.py | 16 ++++++++-------- .../stable_diffusion/test_stable_diffusion.py | 2 +- tests/test_scheduler.py | 7 +++++-- 8 files changed, 41 insertions(+), 28 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 80e1edc4ab4f..964258cbaea6 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -52,9 +52,7 @@ from .utils.dummy_pt_objects import * # noqa F403 if is_torch_available() and is_scipy_available(): - from .schedulers import LMSDiscreteScheduler - from .schedulers import EulerAncestralDiscreteScheduler - from .schedulers import EulerDiscreteScheduler + from .schedulers import EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler else: from .utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 7eac14c39732..b49f2a8b5f10 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -9,8 +9,13 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, \ -EulerDiscreteScheduler, EulerAncestralDiscreteScheduler +from ...schedulers import ( + DDIMScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) from ...utils import deprecate, logging from . import StableDiffusionPipelineOutput from .safety_checker import StableDiffusionSafetyChecker @@ -53,8 +58,9 @@ def __init__( text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, - scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, - EulerDiscreteScheduler, EulerAncestralDiscreteScheduler], + scheduler: Union[ + DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler + ], safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPFeatureExtractor, ): diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index dbf3ed738cf4..1674297aa288 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -10,8 +10,13 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKL, UNet2DConditionModel from ...pipeline_utils import DiffusionPipeline -from ...schedulers import DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, \ -EulerDiscreteScheduler, EulerAncestralDiscreteScheduler +from ...schedulers import ( + DDIMScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) from ...utils import deprecate, logging from . import StableDiffusionPipelineOutput from .safety_checker import StableDiffusionSafetyChecker @@ -65,8 +70,7 @@ def __init__( tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, scheduler: Union[ - DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, - EulerDiscreteScheduler, EulerAncestralDiscreteScheduler + DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler ], safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPFeatureExtractor, diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 6db70be54bb6..b7cb0e127db5 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -41,8 +41,8 @@ if is_scipy_available() and is_torch_available(): - from .scheduling_lms_discrete import LMSDiscreteScheduler from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler from .scheduling_euler_discrete import EulerDiscreteScheduler + from .scheduling_lms_discrete import LMSDiscreteScheduler else: from ..utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index b3b2866bc715..717f262c927e 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -191,8 +191,8 @@ def step( Returns: [`~schedulers.scheduling_utils.EulerAncestralDiscreteSchedulerOutput`] or `tuple`: - [`~schedulers.scheduling_utils.EulerAncestralDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. - When returning a tuple, the first element is the sample tensor. + [`~schedulers.scheduling_utils.EulerAncestralDiscreteSchedulerOutput`] if `return_dict` is True, otherwise + a `tuple`. When returning a tuple, the first element is the sample tensor. """ if not self.is_scale_input_called: @@ -225,8 +225,8 @@ def step( pred_original_sample = sample - sigma * model_output sigma_from = self.sigmas[step_index] sigma_to = self.sigmas[step_index + 1] - sigma_up = (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5 - sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5 + sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5 + sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5 # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma self.derivatives.append(derivative) @@ -241,7 +241,9 @@ def step( if not return_dict: return (prev_sample,) - return EulerAncestralDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) + return EulerAncestralDiscreteSchedulerOutput( + prev_sample=prev_sample, pred_original_sample=pred_original_sample + ) def add_noise( self, diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index a76c31bbd392..cf34d2009ea4 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -175,10 +175,10 @@ def step( timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, order: int = 4, - s_churn: float = 0., - s_tmin: float = 0., - s_tmax: float = float('inf'), - s_noise: float = 1., + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, return_dict: bool = True, ) -> Union[EulerDiscreteSchedulerOutput, Tuple]: """ @@ -199,8 +199,8 @@ def step( Returns: [`~schedulers.scheduling_utils.EulerDiscreteSchedulerOutput`] or `tuple`: - [`~schedulers.scheduling_utils.EulerDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. - When returning a tuple, the first element is the sample tensor. + [`~schedulers.scheduling_utils.EulerDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a + `tuple`. When returning a tuple, the first element is the sample tensor. """ if not self.is_scale_input_called: @@ -228,11 +228,11 @@ def step( else: step_index = (self.timesteps == timestep).nonzero().item() sigma = self.sigmas[step_index] - gamma = min(s_churn / (len(self.sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigma <= s_tmax else 0. + gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 eps = torch.randn_like(sample) * s_noise sigma_hat = sigma * (gamma + 1) if gamma > 0: - sample = sample + eps * (sigma_hat ** 2 - sigma ** 2) ** 0.5 + sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise pred_original_sample = sample - sigma_hat * model_output diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index 610efcd1deb6..f9550676cd66 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -23,9 +23,9 @@ from diffusers import ( AutoencoderKL, DDIMScheduler, - LMSDiscreteScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, + LMSDiscreteScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel, diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index dbfc2290ba45..a03f261a40a0 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -22,10 +22,10 @@ from diffusers import ( DDIMScheduler, DDPMScheduler, - IPNDMScheduler, - LMSDiscreteScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, + IPNDMScheduler, + LMSDiscreteScheduler, PNDMScheduler, ScoreSdeVeScheduler, ) @@ -934,6 +934,7 @@ def test_full_loop_no_noise(self): assert abs(result_sum.item() - 1006.388) < 1e-2 assert abs(result_mean.item() - 1.31) < 1e-3 + class EulerDiscreteSchedulerTest(SchedulerCommonTest): scheduler_classes = (EulerDiscreteScheduler,) num_inference_steps = 10 @@ -990,6 +991,7 @@ def test_full_loop_no_noise(self): assert abs(result_sum.item() - 1006.388) < 1e-2 assert abs(result_mean.item() - 1.31) < 1e-3 + class EulerAncestralDiscreteSchedulerTest(SchedulerCommonTest): scheduler_classes = (EulerAncestralDiscreteScheduler,) num_inference_steps = 10 @@ -1046,6 +1048,7 @@ def test_full_loop_no_noise(self): assert abs(result_sum.item() - 1006.388) < 1e-2 assert abs(result_mean.item() - 1.31) < 1e-3 + class IPNDMSchedulerTest(SchedulerCommonTest): scheduler_classes = (IPNDMScheduler,) forward_default_kwargs = (("num_inference_steps", 50),) From 549664942790d1d20cbed42e49206e6d4fca00d7 Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Thu, 27 Oct 2022 12:50:15 +0100 Subject: [PATCH 03/27] make fix-copies --- .../utils/dummy_torch_and_scipy_objects.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/src/diffusers/utils/dummy_torch_and_scipy_objects.py b/src/diffusers/utils/dummy_torch_and_scipy_objects.py index 13f17349bb45..f1ee08f04c52 100644 --- a/src/diffusers/utils/dummy_torch_and_scipy_objects.py +++ b/src/diffusers/utils/dummy_torch_and_scipy_objects.py @@ -4,6 +4,36 @@ from ..utils import DummyObject, requires_backends +class EulerAncestralDiscreteScheduler(metaclass=DummyObject): + _backends = ["torch", "scipy"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "scipy"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "scipy"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "scipy"]) + + +class EulerDiscreteScheduler(metaclass=DummyObject): + _backends = ["torch", "scipy"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "scipy"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "scipy"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "scipy"]) + + class LMSDiscreteScheduler(metaclass=DummyObject): _backends = ["torch", "scipy"] From 67536449a53d849a56a584f9b8d3f1a92e36c4a6 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Fri, 28 Oct 2022 16:45:26 +0200 Subject: [PATCH 04/27] fix tests for euler a --- tests/test_scheduler.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index a03f261a40a0..3b79c8eb13cf 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -19,6 +19,7 @@ import numpy as np import torch +from black import out from diffusers import ( DDIMScheduler, DDPMScheduler, @@ -75,6 +76,7 @@ def model(sample, t, *args): return model def check_over_configs(self, time_step=0, **config): + print(config) kwargs = dict(self.forward_default_kwargs) num_inference_steps = kwargs.pop("num_inference_steps", None) @@ -96,7 +98,11 @@ def check_over_configs(self, time_step=0, **config): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps + # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler + torch.manual_seed(0) output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample + + torch.manual_seed(0) new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -229,6 +235,8 @@ def recursive_check(tuple_object, dict_object): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps + # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler + torch.manual_seed(0) outputs_dict = scheduler.step(residual, timestep, sample, **kwargs) if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"): @@ -236,6 +244,8 @@ def recursive_check(tuple_object, dict_object): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps + # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler + torch.manual_seed(0) outputs_tuple = scheduler.step(residual, timestep, sample, return_dict=False, **kwargs) recursive_check(outputs_tuple, outputs_dict) @@ -987,9 +997,10 @@ def test_full_loop_no_noise(self): result_sum = torch.sum(torch.abs(sample)) result_mean = torch.mean(torch.abs(sample)) + print(result_sum, result_mean) - assert abs(result_sum.item() - 1006.388) < 1e-2 - assert abs(result_mean.item() - 1.31) < 1e-3 + assert abs(result_sum.item() - 10.0807) < 1e-2 + assert abs(result_mean.item() - 0.0131) < 1e-3 class EulerAncestralDiscreteSchedulerTest(SchedulerCommonTest): @@ -1025,6 +1036,7 @@ def test_time_indices(self): self.check_over_forward(time_step=t) def test_full_loop_no_noise(self): + torch.manual_seed(0) scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() scheduler = scheduler_class(**scheduler_config) @@ -1045,8 +1057,8 @@ def test_full_loop_no_noise(self): result_sum = torch.sum(torch.abs(sample)) result_mean = torch.mean(torch.abs(sample)) - assert abs(result_sum.item() - 1006.388) < 1e-2 - assert abs(result_mean.item() - 1.31) < 1e-3 + assert abs(result_sum.item() - 144.3660) < 1e-2 + assert abs(result_mean.item() - 0.1880) < 1e-3 class IPNDMSchedulerTest(SchedulerCommonTest): From 85ae890e49960a918cede737bcb4911c68d62d50 Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Fri, 28 Oct 2022 16:45:35 +0100 Subject: [PATCH 05/27] Update src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py Co-authored-by: Anton Lozhkov --- src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 717f262c927e..f2ae07a795a3 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -26,7 +26,7 @@ @dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete +# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerAncestralDiscrete class EulerAncestralDiscreteSchedulerOutput(BaseOutput): """ Output class for the scheduler's step function output. From bdc83349dbc9e19d04d47f1dab9e6e7be59cf261 Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Fri, 28 Oct 2022 16:45:48 +0100 Subject: [PATCH 06/27] Update src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py Co-authored-by: Anton Lozhkov --- src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index f2ae07a795a3..0983c5009393 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -109,7 +109,7 @@ def scale_model_input( self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] ) -> torch.FloatTensor: """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. + Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. Args: sample (`torch.FloatTensor`): input sample From 003a61b47e2b0637dd2de7fb817a64343914a20a Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Fri, 28 Oct 2022 16:45:59 +0100 Subject: [PATCH 07/27] Update src/diffusers/schedulers/scheduling_euler_discrete.py Co-authored-by: Anton Lozhkov --- src/diffusers/schedulers/scheduling_euler_discrete.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index cf34d2009ea4..d7d7dcfdae5f 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -109,7 +109,7 @@ def scale_model_input( self, sample: torch.FloatTensor, timestep: Union[float, torch.FloatTensor] ) -> torch.FloatTensor: """ - Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm. + Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm. Args: sample (`torch.FloatTensor`): input sample From 8a3b4a386ab56308d851bf5573147564db2a6a82 Mon Sep 17 00:00:00 2001 From: hlky <106811348+hlky@users.noreply.github.com> Date: Fri, 28 Oct 2022 16:46:14 +0100 Subject: [PATCH 08/27] Update src/diffusers/schedulers/scheduling_euler_discrete.py Co-authored-by: Anton Lozhkov --- src/diffusers/schedulers/scheduling_euler_discrete.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index d7d7dcfdae5f..d755adcad6b0 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -26,7 +26,7 @@ @dataclass -# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->LMSDiscrete +# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete class EulerDiscreteSchedulerOutput(BaseOutput): """ Output class for the scheduler's step function output. From 44c050962b1caf59e99768c601fcf4094f63324d Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Fri, 28 Oct 2022 18:14:51 +0200 Subject: [PATCH 09/27] remove unused arg and method --- .../scheduling_euler_ancestral_discrete.py | 32 ++---------------- .../schedulers/scheduling_euler_discrete.py | 33 ++----------------- 2 files changed, 5 insertions(+), 60 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 0983c5009393..c3bab3e4c518 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -102,7 +102,6 @@ def __init__( self.num_inference_steps = None timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() self.timesteps = torch.from_numpy(timesteps) - self.derivatives = [] self.is_scale_input_called = False def scale_model_input( @@ -126,28 +125,6 @@ def scale_model_input( self.is_scale_input_called = True return sample - def get_lms_coefficient(self, order, t, current_order): - """ - Compute a linear multistep coefficient. - - Args: - order (TODO): - t (TODO): - current_order (TODO): - """ - - def lms_derivative(tau): - prod = 1.0 - for k in range(order): - if current_order == k: - continue - prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) - return prod - - integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] - - return integrated_coeff - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): """ Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -167,14 +144,11 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to(device=device) - self.derivatives = [] - def step( self, model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, - order: int = 4, return_dict: bool = True, ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]: """ @@ -186,7 +160,6 @@ def step( timestep (`float`): current timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. return_dict (`bool`): option for returning tuple rather than EulerAncestralDiscreteSchedulerOutput class Returns: @@ -227,11 +200,10 @@ def step( sigma_to = self.sigmas[step_index + 1] sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5 sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5 + # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma - self.derivatives.append(derivative) - if len(self.derivatives) > order: - self.derivatives.pop(0) + dt = sigma_down - sigma prev_sample = sample + derivative * dt diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index d755adcad6b0..7e63113d0dfb 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -102,7 +102,6 @@ def __init__( self.num_inference_steps = None timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy() self.timesteps = torch.from_numpy(timesteps) - self.derivatives = [] self.is_scale_input_called = False def scale_model_input( @@ -126,28 +125,6 @@ def scale_model_input( self.is_scale_input_called = True return sample - def get_lms_coefficient(self, order, t, current_order): - """ - Compute a linear multistep coefficient. - - Args: - order (TODO): - t (TODO): - current_order (TODO): - """ - - def lms_derivative(tau): - prod = 1.0 - for k in range(order): - if current_order == k: - continue - prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k]) - return prod - - integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0] - - return integrated_coeff - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): """ Sets the timesteps used for the diffusion chain. Supporting function to be run before inference. @@ -167,14 +144,11 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic self.sigmas = torch.from_numpy(sigmas).to(device=device) self.timesteps = torch.from_numpy(timesteps).to(device=device) - self.derivatives = [] - def step( self, model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, - order: int = 4, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float("inf"), @@ -190,7 +164,6 @@ def step( timestep (`float`): current timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. - order: coefficient for multi-step inference. s_churn (`float`) s_tmin (`float`) s_tmax (`float`) @@ -228,19 +201,19 @@ def step( else: step_index = (self.timesteps == timestep).nonzero().item() sigma = self.sigmas[step_index] + gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 eps = torch.randn_like(sample) * s_noise sigma_hat = sigma * (gamma + 1) + if gamma > 0: sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 + # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise pred_original_sample = sample - sigma_hat * model_output # 2. Convert to an ODE derivative derivative = (sample - pred_original_sample) / sigma_hat - self.derivatives.append(derivative) - if len(self.derivatives) > order: - self.derivatives.pop(0) dt = self.sigmas[step_index + 1] - sigma_hat From 0126944db9dc151d8d672c9360b7610b9bfadfb7 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Fri, 28 Oct 2022 18:19:48 +0200 Subject: [PATCH 10/27] update doc --- .../schedulers/scheduling_euler_ancestral_discrete.py | 5 ++--- src/diffusers/schedulers/scheduling_euler_discrete.py | 6 +++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index c3bab3e4c518..126dd0abe28a 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -46,9 +46,8 @@ class EulerAncestralDiscreteSchedulerOutput(BaseOutput): class EulerAncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): """ - Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by - Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 + Ancestral sampling with Euler method steps. Based on the original k-diffusion implementation by Katherine Crowson: + https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L72 [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 7e63113d0dfb..fa2f0b4780ec 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -46,9 +46,9 @@ class EulerDiscreteSchedulerOutput(BaseOutput): class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin): """ - Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by - Katherine Crowson: - https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181 + Euler scheduler (Algorithm 2) from Karras et al. (2022) https://arxiv.org/abs/2206.00364. . Based on the original + k-diffusion implementation by Katherine Crowson: + https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L51 [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`. From deedc4ea061f5b543085c47a1d4708933e8965cb Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Fri, 28 Oct 2022 18:21:32 +0200 Subject: [PATCH 11/27] quality --- src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py | 2 -- src/diffusers/schedulers/scheduling_euler_discrete.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 126dd0abe28a..3b9dabbbe4df 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -18,8 +18,6 @@ import numpy as np import torch -from scipy import integrate - from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput, deprecate from .scheduling_utils import SchedulerMixin diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index fa2f0b4780ec..bf55f3721061 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -18,8 +18,6 @@ import numpy as np import torch -from scipy import integrate - from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput, deprecate from .scheduling_utils import SchedulerMixin From dde3f8dd3dc578d2e71e0894882f70f9e430e39b Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Fri, 28 Oct 2022 18:28:41 +0200 Subject: [PATCH 12/27] make flake happy --- tests/test_scheduler.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 3b79c8eb13cf..37c095116fca 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -19,7 +19,6 @@ import numpy as np import torch -from black import out from diffusers import ( DDIMScheduler, DDPMScheduler, From 2a33db8b408b1fd4340c6c00f45028c3e7c0a1a1 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 11:00:38 +0100 Subject: [PATCH 13/27] use logger instead of warn --- .../schedulers/scheduling_euler_ancestral_discrete.py | 6 ++++-- src/diffusers/schedulers/scheduling_euler_discrete.py | 7 +++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 3b9dabbbe4df..64c9072add1c 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -19,9 +19,11 @@ import torch from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, deprecate +from ..utils import BaseOutput, deprecate, logging from .scheduling_utils import SchedulerMixin +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + @dataclass # Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerAncestralDiscrete @@ -166,7 +168,7 @@ def step( """ if not self.is_scale_input_called: - warnings.warn( + logging.warn( "The `scale_model_input` function should be called before `step` to ensure correct denoising. " "See `StableDiffusionPipeline` for a usage example." ) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index bf55f3721061..3e55506e5049 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -19,10 +19,13 @@ import torch from ..configuration_utils import ConfigMixin, register_to_config -from ..utils import BaseOutput, deprecate +from ..utils import BaseOutput, deprecate, logging from .scheduling_utils import SchedulerMixin +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + @dataclass # Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete class EulerDiscreteSchedulerOutput(BaseOutput): @@ -175,7 +178,7 @@ def step( """ if not self.is_scale_input_called: - warnings.warn( + logger.warn( "The `scale_model_input` function should be called before `step` to ensure correct denoising. " "See `StableDiffusionPipeline` for a usage example." ) From e1d2c8826fcf736f3e6e29c4646f30682267bba8 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 11:06:56 +0100 Subject: [PATCH 14/27] raise error instead of deprication --- .../scheduling_euler_ancestral_discrete.py | 31 +++++++++---------- .../schedulers/scheduling_euler_discrete.py | 30 +++++++++--------- 2 files changed, 29 insertions(+), 32 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 64c9072add1c..3624cbe56623 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import warnings + from dataclasses import dataclass from typing import Optional, Tuple, Union @@ -22,6 +22,7 @@ from ..utils import BaseOutput, deprecate, logging from .scheduling_utils import SchedulerMixin + logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -167,30 +168,28 @@ def step( a `tuple`. When returning a tuple, the first element is the sample tensor. """ - if not self.is_scale_input_called: - logging.warn( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) if ( isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor) ): - deprecate( - "timestep as an index", - "0.8.0", + raise ValueError( "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `EulerAncestralDiscreteScheduler.step()` will not be supported in future versions. Make sure to pass" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" " one of the `scheduler.timesteps` as a timestep.", - standard_warn=False, ) - step_index = timestep - else: - step_index = (self.timesteps == timestep).nonzero().item() + + if not self.is_scale_input_called: + logger.warn( + "The `scale_model_input` function should be called before `step` to ensure correct denoising. " + "See `StableDiffusionPipeline` for a usage example." + ) + + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + + step_index = (self.timesteps == timestep).nonzero().item() sigma = self.sigmas[step_index] # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index 3e55506e5049..be1869cef620 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import warnings + from dataclasses import dataclass from typing import Optional, Tuple, Union @@ -177,30 +177,28 @@ def step( `tuple`. When returning a tuple, the first element is the sample tensor. """ - if not self.is_scale_input_called: - logger.warn( - "The `scale_model_input` function should be called before `step` to ensure correct denoising. " - "See `StableDiffusionPipeline` for a usage example." - ) - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) if ( isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor) ): - deprecate( - "timestep as an index", - "0.8.0", + raise ValueError( "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `EulerDiscreteScheduler.step()` will not be supported in future versions. Make sure to pass" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" " one of the `scheduler.timesteps` as a timestep.", - standard_warn=False, ) - step_index = timestep - else: - step_index = (self.timesteps == timestep).nonzero().item() + + if not self.is_scale_input_called: + logger.warn( + "The `scale_model_input` function should be called before `step` to ensure correct denoising. " + "See `StableDiffusionPipeline` for a usage example." + ) + + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + + step_index = (self.timesteps == timestep).nonzero().item() sigma = self.sigmas[step_index] gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 From 51a855ec424d38c672777d8e3f5db5a4c20ed1e6 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 11:10:21 +0100 Subject: [PATCH 15/27] don't require scipy --- src/diffusers/__init__.py | 4 ++- src/diffusers/schedulers/__init__.py | 4 +-- src/diffusers/utils/dummy_pt_objects.py | 30 +++++++++++++++++++ .../utils/dummy_torch_and_scipy_objects.py | 30 ------------------- 4 files changed, 35 insertions(+), 33 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 964258cbaea6..49c3e82b8e7b 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -41,6 +41,8 @@ from .schedulers import ( DDIMScheduler, DDPMScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, IPNDMScheduler, KarrasVeScheduler, PNDMScheduler, @@ -52,7 +54,7 @@ from .utils.dummy_pt_objects import * # noqa F403 if is_torch_available() and is_scipy_available(): - from .schedulers import EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler + from .schedulers import LMSDiscreteScheduler else: from .utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index b7cb0e127db5..c3999d2cac61 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -19,6 +19,8 @@ if is_torch_available(): from .scheduling_ddim import DDIMScheduler from .scheduling_ddpm import DDPMScheduler + from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler + from .scheduling_euler_discrete import EulerDiscreteScheduler from .scheduling_ipndm import IPNDMScheduler from .scheduling_karras_ve import KarrasVeScheduler from .scheduling_pndm import PNDMScheduler @@ -41,8 +43,6 @@ if is_scipy_available() and is_torch_available(): - from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler - from .scheduling_euler_discrete import EulerDiscreteScheduler from .scheduling_lms_discrete import LMSDiscreteScheduler else: from ..utils.dummy_torch_and_scipy_objects import * # noqa F403 diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 93834c76550b..5dd583279708 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -272,6 +272,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class EulerAncestralDiscreteScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class EulerDiscreteScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class IPNDMScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_scipy_objects.py b/src/diffusers/utils/dummy_torch_and_scipy_objects.py index f1ee08f04c52..13f17349bb45 100644 --- a/src/diffusers/utils/dummy_torch_and_scipy_objects.py +++ b/src/diffusers/utils/dummy_torch_and_scipy_objects.py @@ -4,36 +4,6 @@ from ..utils import DummyObject, requires_backends -class EulerAncestralDiscreteScheduler(metaclass=DummyObject): - _backends = ["torch", "scipy"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "scipy"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) - - -class EulerDiscreteScheduler(metaclass=DummyObject): - _backends = ["torch", "scipy"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "scipy"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "scipy"]) - - class LMSDiscreteScheduler(metaclass=DummyObject): _backends = ["torch", "scipy"] From 18c9d9ade2f62306f0eeefc8a753494310220243 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 11:27:23 +0100 Subject: [PATCH 16/27] pass generator in step --- .../stable_diffusion/pipeline_stable_diffusion.py | 2 +- .../stable_diffusion/pipeline_stable_diffusion_img2img.py | 2 +- .../stable_diffusion/pipeline_stable_diffusion_inpaint.py | 2 +- .../pipeline_stable_diffusion_inpaint_legacy.py | 2 +- .../schedulers/scheduling_euler_ancestral_discrete.py | 6 +++++- src/diffusers/schedulers/scheduling_euler_discrete.py | 7 ++++++- src/diffusers/schedulers/scheduling_ipndm.py | 4 +++- src/diffusers/schedulers/scheduling_karras_ve.py | 2 ++ src/diffusers/schedulers/scheduling_lms_discrete.py | 2 ++ src/diffusers/schedulers/scheduling_pndm.py | 2 ++ 10 files changed, 24 insertions(+), 7 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index b49f2a8b5f10..236cfde24e3b 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -352,7 +352,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index 1674297aa288..d6239b75be77 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -365,7 +365,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py index d49f55aa921d..7e403ed6db2c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -397,7 +397,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py index 71a2be62128c..d6c50f567e57 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py @@ -374,7 +374,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample # masking init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t])) diff --git a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py index 3624cbe56623..b153dfbe2f02 100644 --- a/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py @@ -149,6 +149,7 @@ def step( model_output: torch.FloatTensor, timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]: """ @@ -160,6 +161,7 @@ def step( timestep (`float`): current timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. + generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than EulerAncestralDiscreteSchedulerOutput class Returns: @@ -206,7 +208,9 @@ def step( prev_sample = sample + derivative * dt - prev_sample = prev_sample + torch.randn_like(prev_sample) * sigma_up + device = model_output.device if torch.is_tensor(model_output) else "cpu" + noise = torch.randn(model_output.shape, dtype=model_output.dtype, generator=generator).to(device) + prev_sample = prev_sample + noise * sigma_up if not return_dict: return (prev_sample,) diff --git a/src/diffusers/schedulers/scheduling_euler_discrete.py b/src/diffusers/schedulers/scheduling_euler_discrete.py index be1869cef620..a71fbda13df6 100644 --- a/src/diffusers/schedulers/scheduling_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_euler_discrete.py @@ -154,6 +154,7 @@ def step( s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[EulerDiscreteSchedulerOutput, Tuple]: """ @@ -169,6 +170,7 @@ def step( s_tmin (`float`) s_tmax (`float`) s_noise (`float`) + generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than EulerDiscreteSchedulerOutput class Returns: @@ -202,7 +204,10 @@ def step( sigma = self.sigmas[step_index] gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0 - eps = torch.randn_like(sample) * s_noise + + device = model_output.device if torch.is_tensor(model_output) else "cpu" + noise = torch.randn(model_output.shape, dtype=model_output.dtype, generator=generator).to(device) + eps = noise * s_noise sigma_hat = sigma * (gamma + 1) if gamma > 0: diff --git a/src/diffusers/schedulers/scheduling_ipndm.py b/src/diffusers/schedulers/scheduling_ipndm.py index fb413a280530..1920230b0d2e 100644 --- a/src/diffusers/schedulers/scheduling_ipndm.py +++ b/src/diffusers/schedulers/scheduling_ipndm.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Tuple, Union +from typing import Optional, Tuple, Union import torch @@ -78,6 +78,7 @@ def step( model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[SchedulerOutput, Tuple]: """ @@ -89,6 +90,7 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. + generator (`torch.Generator`, `optional`): Random number generator. return_dict (`bool`): option for returning tuple rather than SchedulerOutput class Returns: diff --git a/src/diffusers/schedulers/scheduling_karras_ve.py b/src/diffusers/schedulers/scheduling_karras_ve.py index 743f2e061c53..b7bbbf349ac1 100644 --- a/src/diffusers/schedulers/scheduling_karras_ve.py +++ b/src/diffusers/schedulers/scheduling_karras_ve.py @@ -157,6 +157,7 @@ def step( sigma_hat: float, sigma_prev: float, sample_hat: torch.FloatTensor, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[KarrasVeOutput, Tuple]: """ @@ -168,6 +169,7 @@ def step( sigma_hat (`float`): TODO sigma_prev (`float`): TODO sample_hat (`torch.FloatTensor`): TODO + generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than KarrasVeOutput class KarrasVeOutput: updated sample in the diffusion chain and derivative (TODO double check). diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 6157f4b4dc65..fad5a992abd1 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -175,6 +175,7 @@ def step( timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, order: int = 4, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[LMSDiscreteSchedulerOutput, Tuple]: """ @@ -187,6 +188,7 @@ def step( sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. order: coefficient for multi-step inference. + generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than LMSDiscreteSchedulerOutput class Returns: diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py index 03ea5c15d1b0..39e5ddba96c3 100644 --- a/src/diffusers/schedulers/scheduling_pndm.py +++ b/src/diffusers/schedulers/scheduling_pndm.py @@ -185,6 +185,7 @@ def step( model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[SchedulerOutput, Tuple]: """ @@ -198,6 +199,7 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. + generator (`torch.Generator`, `optional`): Random number generator. return_dict (`bool`): option for returning tuple rather than SchedulerOutput class Returns: From 66ee52e20b20a7f0fa0754af53256fe648381ab4 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 11:50:36 +0100 Subject: [PATCH 17/27] fix tests --- tests/test_scheduler.py | 83 +++++++++++++++++++++++++---------------- 1 file changed, 51 insertions(+), 32 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 37c095116fca..cdf7b3fe4b49 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -81,6 +81,9 @@ def check_over_configs(self, time_step=0, **config): num_inference_steps = kwargs.pop("num_inference_steps", None) for scheduler_class in self.scheduler_classes: + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + time_step = float(time_step) + sample = self.dummy_sample residual = 0.1 * sample @@ -98,11 +101,11 @@ def check_over_configs(self, time_step=0, **config): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - torch.manual_seed(0) - output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample + generator = torch.Generator().manual_seed(0) + output = scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample - torch.manual_seed(0) - new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample + generator = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -113,6 +116,9 @@ def check_over_forward(self, time_step=0, **forward_kwargs): num_inference_steps = kwargs.pop("num_inference_steps", None) for scheduler_class in self.scheduler_classes: + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + time_step = float(time_step) + sample = self.dummy_sample residual = 0.1 * sample @@ -129,10 +135,11 @@ def check_over_forward(self, time_step=0, **forward_kwargs): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - torch.manual_seed(0) - output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - torch.manual_seed(0) - new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample + generator = torch.Generator().manual_seed(0) + output = scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample + + generator = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -148,6 +155,10 @@ def test_from_pretrained_save_pretrained(self): scheduler_config = self.get_scheduler_config() scheduler = scheduler_class(**scheduler_config) + timestep = 1 + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + timestep = float(timestep) + with tempfile.TemporaryDirectory() as tmpdirname: scheduler.save_config(tmpdirname) new_scheduler = scheduler_class.from_config(tmpdirname) @@ -158,10 +169,11 @@ def test_from_pretrained_save_pretrained(self): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - torch.manual_seed(0) - output = scheduler.step(residual, 1, sample, **kwargs).prev_sample - torch.manual_seed(0) - new_output = new_scheduler.step(residual, 1, sample, **kwargs).prev_sample + generator = torch.Generator().manual_seed(0) + output = scheduler.step(residual, timestep, sample, generator=generator, **kwargs).prev_sample + + generator = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, timestep, sample, generator=generator, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -170,7 +182,14 @@ def test_step_shape(self): num_inference_steps = kwargs.pop("num_inference_steps", None) + timestep_0 = 0 + timestep_1 = 1 + for scheduler_class in self.scheduler_classes: + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + timestep_0 = float(timestep_0) + timestep_1 = float(timestep_1) + scheduler_config = self.get_scheduler_config() scheduler = scheduler_class(**scheduler_config) @@ -182,8 +201,8 @@ def test_step_shape(self): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - output_0 = scheduler.step(residual, 0, sample, **kwargs).prev_sample - output_1 = scheduler.step(residual, 1, sample, **kwargs).prev_sample + output_0 = scheduler.step(residual, timestep_0, sample, **kwargs).prev_sample + output_1 = scheduler.step(residual, timestep_1, sample, **kwargs).prev_sample self.assertEqual(output_0.shape, sample.shape) self.assertEqual(output_0.shape, output_1.shape) @@ -223,6 +242,9 @@ def recursive_check(tuple_object, dict_object): timestep = 1 for scheduler_class in self.scheduler_classes: + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + timestep = float(timestep) + scheduler_config = self.get_scheduler_config() scheduler = scheduler_class(**scheduler_config) @@ -235,8 +257,8 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - torch.manual_seed(0) - outputs_dict = scheduler.step(residual, timestep, sample, **kwargs) + generator = torch.Generator().manual_seed(0) + outputs_dict = scheduler.step(residual, timestep, sample, generator=generator, **kwargs) if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"): scheduler.set_timesteps(num_inference_steps) @@ -244,8 +266,10 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - torch.manual_seed(0) - outputs_tuple = scheduler.step(residual, timestep, sample, return_dict=False, **kwargs) + generator = torch.Generator().manual_seed(0) + outputs_tuple = scheduler.step( + residual, timestep, sample, generator=generator, return_dict=False, **kwargs + ) recursive_check(outputs_tuple, outputs_dict) @@ -972,10 +996,6 @@ def test_schedules(self): for schedule in ["linear", "scaled_linear"]: self.check_over_configs(beta_schedule=schedule) - def test_time_indices(self): - for t in [0, 500, 800]: - self.check_over_forward(time_step=t) - def test_full_loop_no_noise(self): scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() @@ -983,6 +1003,8 @@ def test_full_loop_no_noise(self): scheduler.set_timesteps(self.num_inference_steps) + generator = torch.Generator().manual_seed(0) + model = self.dummy_model() sample = self.dummy_sample_deter * scheduler.init_noise_sigma @@ -991,7 +1013,7 @@ def test_full_loop_no_noise(self): model_output = model(sample, t) - output = scheduler.step(model_output, t, sample) + output = scheduler.step(model_output, t, sample, generator=generator) sample = output.prev_sample result_sum = torch.sum(torch.abs(sample)) @@ -1030,18 +1052,15 @@ def test_schedules(self): for schedule in ["linear", "scaled_linear"]: self.check_over_configs(beta_schedule=schedule) - def test_time_indices(self): - for t in [0, 500, 800]: - self.check_over_forward(time_step=t) - def test_full_loop_no_noise(self): - torch.manual_seed(0) scheduler_class = self.scheduler_classes[0] scheduler_config = self.get_scheduler_config() scheduler = scheduler_class(**scheduler_config) scheduler.set_timesteps(self.num_inference_steps) + generator = torch.Generator().manual_seed(0) + model = self.dummy_model() sample = self.dummy_sample_deter * scheduler.init_noise_sigma @@ -1050,14 +1069,14 @@ def test_full_loop_no_noise(self): model_output = model(sample, t) - output = scheduler.step(model_output, t, sample) + output = scheduler.step(model_output, t, sample, generator=generator) sample = output.prev_sample result_sum = torch.sum(torch.abs(sample)) result_mean = torch.mean(torch.abs(sample)) - - assert abs(result_sum.item() - 144.3660) < 1e-2 - assert abs(result_mean.item() - 0.1880) < 1e-3 + print(result_sum, result_mean) + assert abs(result_sum.item() - 152.3192) < 1e-2 + assert abs(result_mean.item() - 0.1983) < 1e-3 class IPNDMSchedulerTest(SchedulerCommonTest): From 3198b777ffec37e8e38b897dedf8075ada7a967d Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Mon, 31 Oct 2022 13:37:23 +0100 Subject: [PATCH 18/27] Apply suggestions from code review Co-authored-by: Pedro Cuenca --- tests/test_scheduler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index cdf7b3fe4b49..71ee742d332f 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -75,7 +75,6 @@ def model(sample, t, *args): return model def check_over_configs(self, time_step=0, **config): - print(config) kwargs = dict(self.forward_default_kwargs) num_inference_steps = kwargs.pop("num_inference_steps", None) @@ -100,7 +99,7 @@ def check_over_configs(self, time_step=0, **config): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler + # Set the seed before step() as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler generator = torch.Generator().manual_seed(0) output = scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample From 30db08a8c9b241465ee1aab3e037cb2bb1661ad1 Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Mon, 31 Oct 2022 14:04:16 +0100 Subject: [PATCH 19/27] Update tests/test_scheduler.py Co-authored-by: Patrick von Platen --- tests/test_scheduler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 71ee742d332f..33c8fd6542fa 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -79,6 +79,7 @@ def check_over_configs(self, time_step=0, **config): num_inference_steps = kwargs.pop("num_inference_steps", None) + # TODO(Suraj) - delete the following two lines once DDPM, DDIM, and PNDM have timesteps casted to float by default for scheduler_class in self.scheduler_classes: if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): time_step = float(time_step) From 9cf1cf066a06f9aab89a2238bbbbb4a2fe0e2884 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:09:39 +0100 Subject: [PATCH 20/27] remove unused generator --- src/diffusers/schedulers/scheduling_ipndm.py | 2 -- src/diffusers/schedulers/scheduling_karras_ve.py | 2 -- src/diffusers/schedulers/scheduling_lms_discrete.py | 2 -- src/diffusers/schedulers/scheduling_pndm.py | 2 -- 4 files changed, 8 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ipndm.py b/src/diffusers/schedulers/scheduling_ipndm.py index 1920230b0d2e..71455d573bab 100644 --- a/src/diffusers/schedulers/scheduling_ipndm.py +++ b/src/diffusers/schedulers/scheduling_ipndm.py @@ -78,7 +78,6 @@ def step( model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, - generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[SchedulerOutput, Tuple]: """ @@ -90,7 +89,6 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. - generator (`torch.Generator`, `optional`): Random number generator. return_dict (`bool`): option for returning tuple rather than SchedulerOutput class Returns: diff --git a/src/diffusers/schedulers/scheduling_karras_ve.py b/src/diffusers/schedulers/scheduling_karras_ve.py index b7bbbf349ac1..743f2e061c53 100644 --- a/src/diffusers/schedulers/scheduling_karras_ve.py +++ b/src/diffusers/schedulers/scheduling_karras_ve.py @@ -157,7 +157,6 @@ def step( sigma_hat: float, sigma_prev: float, sample_hat: torch.FloatTensor, - generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[KarrasVeOutput, Tuple]: """ @@ -169,7 +168,6 @@ def step( sigma_hat (`float`): TODO sigma_prev (`float`): TODO sample_hat (`torch.FloatTensor`): TODO - generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than KarrasVeOutput class KarrasVeOutput: updated sample in the diffusion chain and derivative (TODO double check). diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index fad5a992abd1..6157f4b4dc65 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -175,7 +175,6 @@ def step( timestep: Union[float, torch.FloatTensor], sample: torch.FloatTensor, order: int = 4, - generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[LMSDiscreteSchedulerOutput, Tuple]: """ @@ -188,7 +187,6 @@ def step( sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. order: coefficient for multi-step inference. - generator (`torch.Generator`, optional): Random number generator. return_dict (`bool`): option for returning tuple rather than LMSDiscreteSchedulerOutput class Returns: diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py index 39e5ddba96c3..03ea5c15d1b0 100644 --- a/src/diffusers/schedulers/scheduling_pndm.py +++ b/src/diffusers/schedulers/scheduling_pndm.py @@ -185,7 +185,6 @@ def step( model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor, - generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[SchedulerOutput, Tuple]: """ @@ -199,7 +198,6 @@ def step( timestep (`int`): current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): current instance of sample being created by diffusion process. - generator (`torch.Generator`, `optional`): Random number generator. return_dict (`bool`): option for returning tuple rather than SchedulerOutput class Returns: From b1324caa2d46977affd8827a54333e40f49bc58b Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:11:20 +0100 Subject: [PATCH 21/27] pass generator as extra_step_kwargs --- .../stable_diffusion/pipeline_stable_diffusion.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 236cfde24e3b..28bdab1ac6c5 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -338,6 +338,11 @@ def __call__( if accepts_eta: extra_step_kwargs["eta"] = eta + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + for i, t in enumerate(self.progress_bar(timesteps_tensor)): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents @@ -352,7 +357,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: From c7fe0a0993a1f65e8f8db29b1dfc9ec21d1ea87a Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:19:56 +0100 Subject: [PATCH 22/27] update tests --- tests/test_scheduler.py | 42 +++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 33c8fd6542fa..1602a998cde6 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -101,11 +101,13 @@ def check_over_configs(self, time_step=0, **config): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before step() as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - generator = torch.Generator().manual_seed(0) - output = scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - generator = torch.Generator().manual_seed(0) - new_output = new_scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -135,11 +137,13 @@ def check_over_forward(self, time_step=0, **forward_kwargs): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - generator = torch.Generator().manual_seed(0) - output = scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - generator = torch.Generator().manual_seed(0) - new_output = new_scheduler.step(residual, time_step, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -169,11 +173,13 @@ def test_from_pretrained_save_pretrained(self): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - generator = torch.Generator().manual_seed(0) - output = scheduler.step(residual, timestep, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + output = scheduler.step(residual, timestep, sample, **kwargs).prev_sample - generator = torch.Generator().manual_seed(0) - new_output = new_scheduler.step(residual, timestep, sample, generator=generator, **kwargs).prev_sample + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + new_output = new_scheduler.step(residual, timestep, sample, **kwargs).prev_sample assert torch.sum(torch.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical" @@ -257,8 +263,9 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - generator = torch.Generator().manual_seed(0) - outputs_dict = scheduler.step(residual, timestep, sample, generator=generator, **kwargs) + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + outputs_dict = scheduler.step(residual, timestep, sample, **kwargs) if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"): scheduler.set_timesteps(num_inference_steps) @@ -266,10 +273,9 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - generator = torch.Generator().manual_seed(0) - outputs_tuple = scheduler.step( - residual, timestep, sample, generator=generator, return_dict=False, **kwargs - ) + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + kwargs["generator"] = torch.Generator().manual_seed(0) + outputs_tuple = scheduler.step(residual, timestep, sample, return_dict=False, **kwargs) recursive_check(outputs_tuple, outputs_dict) From c5e6aa5cf19952226994ad252744c51fcb21f1e4 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:23:44 +0100 Subject: [PATCH 23/27] pass generator as kwarg --- .../stable_diffusion/pipeline_stable_diffusion_img2img.py | 7 ++++++- .../stable_diffusion/pipeline_stable_diffusion_inpaint.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index d6239b75be77..799b1c19f6c1 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -343,6 +343,11 @@ def __call__( if accepts_eta: extra_step_kwargs["eta"] = eta + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + latents = init_latents t_start = max(num_inference_steps - init_timestep + offset, 0) @@ -365,7 +370,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py index 7e403ed6db2c..77bd55f2744d 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -379,6 +379,11 @@ def __call__( if accepts_eta: extra_step_kwargs["eta"] = eta + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + for i, t in enumerate(self.progress_bar(timesteps_tensor)): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents @@ -397,7 +402,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided if callback is not None and i % callback_steps == 0: From d6daae7ef2b37e174fab030844233856eafc3296 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:25:26 +0100 Subject: [PATCH 24/27] pass generator as kwarg --- .../pipeline_stable_diffusion_inpaint_legacy.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py index d6c50f567e57..55472f34cf68 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint_legacy.py @@ -352,6 +352,11 @@ def __call__( if accepts_eta: extra_step_kwargs["eta"] = eta + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + latents = init_latents t_start = max(num_inference_steps - init_timestep + offset, 0) @@ -374,7 +379,7 @@ def __call__( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, generator=generator, **extra_step_kwargs).prev_sample + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # masking init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t])) From 5993631e2e5d18a0d94a47d379db44f389b8bd33 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:25:31 +0100 Subject: [PATCH 25/27] quality --- src/diffusers/schedulers/scheduling_ipndm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_ipndm.py b/src/diffusers/schedulers/scheduling_ipndm.py index 71455d573bab..fb413a280530 100644 --- a/src/diffusers/schedulers/scheduling_ipndm.py +++ b/src/diffusers/schedulers/scheduling_ipndm.py @@ -13,7 +13,7 @@ # limitations under the License. import math -from typing import Optional, Tuple, Union +from typing import Tuple, Union import torch From 6d484c354f4d8b59aa05a9da32a550570d1c5e36 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:33:36 +0100 Subject: [PATCH 26/27] fix test for lms --- tests/test_scheduler.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 1602a998cde6..b70290d001b3 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -101,11 +101,11 @@ def check_over_configs(self, time_step=0, **config): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before step() as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample @@ -137,11 +137,11 @@ def check_over_forward(self, time_step=0, **forward_kwargs): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample @@ -173,11 +173,11 @@ def test_from_pretrained_save_pretrained(self): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, timestep, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, timestep, sample, **kwargs).prev_sample @@ -263,7 +263,7 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) outputs_dict = scheduler.step(residual, timestep, sample, **kwargs) @@ -273,7 +273,7 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LMSDiscreteScheduler): + if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): kwargs["generator"] = torch.Generator().manual_seed(0) outputs_tuple = scheduler.step(residual, timestep, sample, return_dict=False, **kwargs) From 207a5d2cb8fafaf698082369efb13f7e0f478ad5 Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 31 Oct 2022 14:40:12 +0100 Subject: [PATCH 27/27] fix tests --- tests/test_scheduler.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index b70290d001b3..9285eed20ff0 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import inspect import tempfile import unittest from typing import Dict, List, Tuple @@ -101,11 +102,11 @@ def check_over_configs(self, time_step=0, **config): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before step() as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample @@ -137,11 +138,11 @@ def check_over_forward(self, time_step=0, **forward_kwargs): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, time_step, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, time_step, sample, **kwargs).prev_sample @@ -173,11 +174,11 @@ def test_from_pretrained_save_pretrained(self): elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"): kwargs["num_inference_steps"] = num_inference_steps - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) output = scheduler.step(residual, timestep, sample, **kwargs).prev_sample - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) new_output = new_scheduler.step(residual, timestep, sample, **kwargs).prev_sample @@ -263,7 +264,7 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) outputs_dict = scheduler.step(residual, timestep, sample, **kwargs) @@ -273,7 +274,7 @@ def recursive_check(tuple_object, dict_object): kwargs["num_inference_steps"] = num_inference_steps # Set the seed before state as some schedulers are stochastic like EulerAncestralDiscreteScheduler, EulerDiscreteScheduler - if scheduler_class in (EulerAncestralDiscreteScheduler, EulerDiscreteScheduler): + if "generator" in set(inspect.signature(scheduler.step).parameters.keys()): kwargs["generator"] = torch.Generator().manual_seed(0) outputs_tuple = scheduler.step(residual, timestep, sample, return_dict=False, **kwargs)