From a5512d49d50bcf6291fd747a8812c159bc4fe730 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 19 Oct 2023 05:39:54 +0000 Subject: [PATCH] move to the end of loop --- .../pipeline_stable_diffusion.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 9209268b87..2d78fc3e03 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -729,19 +729,6 @@ class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Lo with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_inputs: - callback_kwargs[k] = locals()[k] - callback_kwargs = callback_on_step_end(i, t, callback_kwargs) - - latents = callback_kwargs.pop("latents", latents) - guidance_scale = callback_kwargs.pop("guidance_scale", guidance_scale) - prompt_embeds = callback_kwargs.pop("prompt_embeds", prompt_embeds) - cross_attention_kwargs = callback_kwargs.pop("cross_attention_kwargs", cross_attention_kwargs) - guidance_rescale = callback_kwargs.pop("guidance_rescale", guidance_rescale) - do_classifier_free_guidance = guidance_scale > 1.0 - # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) @@ -767,6 +754,19 @@ class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Lo # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_inputs: + callback_kwargs[k] = locals()[k] + callback_kwargs = callback_on_step_end(i, t, callback_kwargs) + + latents = callback_kwargs.pop("latents", latents) + guidance_scale = callback_kwargs.pop("guidance_scale", guidance_scale) + prompt_embeds = callback_kwargs.pop("prompt_embeds", prompt_embeds) + cross_attention_kwargs = callback_kwargs.pop("cross_attention_kwargs", cross_attention_kwargs) + guidance_rescale = callback_kwargs.pop("guidance_rescale", guidance_rescale) + do_classifier_free_guidance = guidance_scale > 1.0 + # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update()