diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b29f8794b7..5f57027955 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -45,11 +45,11 @@ if is_torch_available(): DDPMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler, - KDPM2DiscreteScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler, IPNDMScheduler, KarrasVeScheduler, + KDPM2DiscreteScheduler, PNDMScheduler, RePaintScheduler, SchedulerMixin, diff --git a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py index 9146d45bd3..70744e0917 100644 --- a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py +++ b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion.py @@ -576,7 +576,7 @@ class AltDiffusionPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py index 16dbd626cd..142d5c6ae3 100644 --- a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py @@ -593,7 +593,7 @@ class AltDiffusionImg2ImgPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py index 9ebbc249f6..6c0b70fe32 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py @@ -681,7 +681,7 @@ class CycleDiffusionPipeline(DiffusionPipeline): ).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index afaef6f481..4bf8e897f8 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -575,7 +575,7 @@ class StableDiffusionPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py index e64a572a87..4c8db74f65 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py @@ -460,7 +460,7 @@ class StableDiffusionImageVariationPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) 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 a25acc0bd1..51b60ee9ff 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -602,7 +602,7 @@ class StableDiffusionImg2ImgPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) 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 6cb2766bc2..3bb0c8cb33 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -721,7 +721,7 @@ class StableDiffusionInpaintPipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) 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 2440b6d5ad..6ead0569fd 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 @@ -617,7 +617,7 @@ class StableDiffusionInpaintPipelineLegacy(DiffusionPipeline): latents = (init_latents_proper * mask) + (latents * (1 - mask)) # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index c9c238ce9a..9d350ff58e 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -535,7 +535,7 @@ class StableDiffusionUpscalePipeline(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py index 7f08e40103..9f5807071c 100644 --- a/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +++ b/src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py @@ -729,7 +729,7 @@ class StableDiffusionPipelineSafe(DiffusionPipeline): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample # call the callback, if provided - if (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0: + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 04544570fa..90a51fe831 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -23,8 +23,8 @@ if is_torch_available(): from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler from .scheduling_euler_discrete import EulerDiscreteScheduler from .scheduling_heun import HeunDiscreteScheduler - from .scheduling_k_dpm_2 import KDPM2DiscreteScheduler from .scheduling_ipndm import IPNDMScheduler + from .scheduling_k_dpm_2 import KDPM2DiscreteScheduler from .scheduling_karras_ve import KarrasVeScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_repaint import RePaintScheduler diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2.py b/src/diffusers/schedulers/scheduling_k_dpm_2.py index 19b76cebf8..d72e05e98d 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2.py @@ -24,7 +24,8 @@ from .scheduling_utils import SchedulerMixin, SchedulerOutput class KDPM2DiscreteScheduler(SchedulerMixin, ConfigMixin): """ - Scheduler created by @crowsonkb in [k_diffusion](https://github.com/crowsonkb/k-diffusion), see: https://github.com/crowsonkb/k-diffusion/blob/5b3af030dd83e0297272d861c19477735d0317ec/k_diffusion/sampling.py#L188 + Scheduler created by @crowsonkb in [k_diffusion](https://github.com/crowsonkb/k-diffusion), see: + https://github.com/crowsonkb/k-diffusion/blob/5b3af030dd83e0297272d861c19477735d0317ec/k_diffusion/sampling.py#L188 Scheduler inspired by DPM-Solver-2 and Algorthim 2 from Karras et al. (2022). @@ -138,7 +139,9 @@ class KDPM2DiscreteScheduler(SchedulerMixin, ConfigMixin): sigmas_interpol = sigmas.log().lerp(sigmas.roll(1).log(), 0.5).exp() self.sigmas = torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2), sigmas[-1:]]) - self.sigmas_interpol = torch.cat([sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2), sigmas_interpol[-1:]]) + self.sigmas_interpol = torch.cat( + [sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2), sigmas_interpol[-1:]] + ) # standard deviation of the initial noise distribution self.init_noise_sigma = self.sigmas.max() diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 9846927cb1..fc937d0d3c 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -407,6 +407,21 @@ class KarrasVeScheduler(metaclass=DummyObject): requires_backends(cls, ["torch"]) +class KDPM2DiscreteScheduler(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 PNDMScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index d49bd4114c..c6922b1206 100755 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -31,8 +31,8 @@ from diffusers import ( EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, HeunDiscreteScheduler, - KDPM2DiscreteScheduler, IPNDMScheduler, + KDPM2DiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler, ScoreSdeVeScheduler,