From 0249d03a0aa1295ae4cd37a45b47fd5ccd8c9229 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Mon, 21 Nov 2022 17:00:03 +0100 Subject: [PATCH] Remove previous denoised prediction from callback. --- .../pipelines/stable_diffusion/pipeline_stable_diffusion.py | 4 ++-- src/diffusers/schedulers/scheduling_heun.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 004379aaf7..81e34b3b83 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -511,11 +511,11 @@ class StableDiffusionPipeline(DiffusionPipeline): noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 - latents, t, pred_original_sample = self.scheduler.step(noise_pred, t, latents, return_dict=False, **extra_step_kwargs) + latents, t = self.scheduler.step(noise_pred, t, latents, return_dict=False, **extra_step_kwargs) # call the callback, if provided if callback is not None and i % callback_steps == 0: - callback(i, t, latents, pred_original_sample) + callback(i, t, latents) i += 1 diff --git a/src/diffusers/schedulers/scheduling_heun.py b/src/diffusers/schedulers/scheduling_heun.py index c704bf8adb..1d77540353 100644 --- a/src/diffusers/schedulers/scheduling_heun.py +++ b/src/diffusers/schedulers/scheduling_heun.py @@ -194,7 +194,7 @@ class HeunDiscreteScheduler(SchedulerMixin, ConfigMixin): print(f"step_index: {step_index}, state_in_first_order: {self.state_in_first_order}, sigma: {sigma}, sigma_next: {sigma_next}, sigma_hat: {sigma_hat}, dt: {dt}") if not return_dict: - return (prev_sample, self.timesteps[step_index], pred_original_sample) + return (prev_sample, self.timesteps[step_index]) return SchedulerOutput(prev_sample=prev_sample, timestep=self.timesteps[step_index])