Unverified Commit bdd16116 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Schedulers] Fix callback steps (#5261)

* fix all

* make fix copies

* make fix copies
parent c8b0f0eb
...@@ -562,7 +562,8 @@ class ComposableStableDiffusionPipeline(DiffusionPipeline): ...@@ -562,7 +562,8 @@ class ComposableStableDiffusionPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# 8. Post-processing # 8. Post-processing
image = self.decode_latents(latents) image = self.decode_latents(latents)
......
...@@ -434,7 +434,8 @@ class ImageToImageInpaintingPipeline(DiffusionPipeline): ...@@ -434,7 +434,8 @@ class ImageToImageInpaintingPipeline(DiffusionPipeline):
# call the callback, if provided # call the callback, if provided
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
latents = 1 / 0.18215 * latents latents = 1 / 0.18215 * latents
image = self.vae.decode(latents).sample image = self.vae.decode(latents).sample
......
...@@ -372,7 +372,8 @@ class StableDiffusionWalkPipeline(DiffusionPipeline): ...@@ -372,7 +372,8 @@ class StableDiffusionWalkPipeline(DiffusionPipeline):
# call the callback, if provided # call the callback, if provided
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
latents = 1 / 0.18215 * latents latents = 1 / 0.18215 * latents
image = self.vae.decode(latents).sample image = self.vae.decode(latents).sample
......
...@@ -1088,7 +1088,8 @@ class StableDiffusionLongPromptWeightingPipeline( ...@@ -1088,7 +1088,8 @@ class StableDiffusionLongPromptWeightingPipeline(
progress_bar.update() progress_bar.update()
if i % callback_steps == 0: if i % callback_steps == 0:
if callback is not None: if callback is not None:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if is_cancelled_callback is not None and is_cancelled_callback(): if is_cancelled_callback is not None and is_cancelled_callback():
return None return None
......
...@@ -846,7 +846,8 @@ class OnnxStableDiffusionLongPromptWeightingPipeline(OnnxStableDiffusionPipeline ...@@ -846,7 +846,8 @@ class OnnxStableDiffusionLongPromptWeightingPipeline(OnnxStableDiffusionPipeline
# call the callback, if provided # call the callback, if provided
if i % callback_steps == 0: if i % callback_steps == 0:
if callback is not None: if callback is not None:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if is_cancelled_callback is not None and is_cancelled_callback(): if is_cancelled_callback is not None and is_cancelled_callback():
return None return None
......
...@@ -1182,7 +1182,8 @@ class SDXLLongPromptWeightingPipeline(DiffusionPipeline, FromSingleFileMixin, Lo ...@@ -1182,7 +1182,8 @@ class SDXLLongPromptWeightingPipeline(DiffusionPipeline, FromSingleFileMixin, Lo
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent": if not output_type == "latent":
# make sure the VAE is in float32 mode, as it overflows in float16 # make sure the VAE is in float32 mode, as it overflows in float16
......
...@@ -202,7 +202,8 @@ class MaskedStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline): ...@@ -202,7 +202,8 @@ class MaskedStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent": if not output_type == "latent":
scaled = latents / self.vae.config.scaling_factor scaled = latents / self.vae.config.scaling_factor
......
...@@ -407,7 +407,8 @@ class MultilingualStableDiffusion(DiffusionPipeline): ...@@ -407,7 +407,8 @@ class MultilingualStableDiffusion(DiffusionPipeline):
# call the callback, if provided # call the callback, if provided
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
latents = 1 / 0.18215 * latents latents = 1 / 0.18215 * latents
image = self.vae.decode(latents).sample image = self.vae.decode(latents).sample
......
...@@ -254,7 +254,8 @@ class Prompt2PromptPipeline(StableDiffusionPipeline): ...@@ -254,7 +254,8 @@ class Prompt2PromptPipeline(StableDiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# 8. Post-processing # 8. Post-processing
if not output_type == "latent": if not output_type == "latent":
......
...@@ -865,7 +865,8 @@ class Zero1to3StableDiffusionPipeline(DiffusionPipeline): ...@@ -865,7 +865,8 @@ class Zero1to3StableDiffusionPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# 8. Post-processing # 8. Post-processing
has_nsfw_concept = None has_nsfw_concept = None
......
...@@ -815,7 +815,8 @@ class OnnxStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline): ...@@ -815,7 +815,8 @@ class OnnxStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent": if not output_type == "latent":
_latents = latents.cpu().detach().numpy() / 0.18215 _latents = latents.cpu().detach().numpy() / 0.18215
......
...@@ -919,7 +919,8 @@ class TensorRTStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline): ...@@ -919,7 +919,8 @@ class TensorRTStableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent": if not output_type == "latent":
_latents = latents.cpu().detach().numpy() / 0.18215 _latents = latents.cpu().detach().numpy() / 0.18215
......
...@@ -337,7 +337,8 @@ class SeedResizeStableDiffusionPipeline(DiffusionPipeline): ...@@ -337,7 +337,8 @@ class SeedResizeStableDiffusionPipeline(DiffusionPipeline):
# call the callback, if provided # call the callback, if provided
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
latents = 1 / 0.18215 * latents latents = 1 / 0.18215 * latents
image = self.vae.decode(latents).sample image = self.vae.decode(latents).sample
......
...@@ -242,7 +242,8 @@ class SpeechToImagePipeline(DiffusionPipeline): ...@@ -242,7 +242,8 @@ class SpeechToImagePipeline(DiffusionPipeline):
# call the callback, if provided # call the callback, if provided
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
latents = 1 / 0.18215 * latents latents = 1 / 0.18215 * latents
image = self.vae.decode(latents).sample image = self.vae.decode(latents).sample
......
...@@ -951,7 +951,8 @@ class StableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline): ...@@ -951,7 +951,8 @@ class StableDiffusionControlNetImg2ImgPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# If we do sequential model offloading, let's offload unet and controlnet # If we do sequential model offloading, let's offload unet and controlnet
# manually for max memory savings # manually for max memory savings
......
...@@ -1100,7 +1100,8 @@ class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline): ...@@ -1100,7 +1100,8 @@ class StableDiffusionControlNetInpaintPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# If we do sequential model offloading, let's offload unet and controlnet # If we do sequential model offloading, let's offload unet and controlnet
# manually for max memory savings # manually for max memory savings
......
...@@ -1081,7 +1081,8 @@ class StableDiffusionControlNetInpaintImg2ImgPipeline(DiffusionPipeline): ...@@ -1081,7 +1081,8 @@ class StableDiffusionControlNetInpaintImg2ImgPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# If we do sequential model offloading, let's offload unet and controlnet # If we do sequential model offloading, let's offload unet and controlnet
# manually for max memory savings # manually for max memory savings
......
...@@ -802,7 +802,8 @@ class StableDiffusionControlNetReferencePipeline(StableDiffusionControlNetPipeli ...@@ -802,7 +802,8 @@ class StableDiffusionControlNetReferencePipeline(StableDiffusionControlNetPipeli
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# If we do sequential model offloading, let's offload unet and controlnet # If we do sequential model offloading, let's offload unet and controlnet
# manually for max memory savings # manually for max memory savings
......
...@@ -817,7 +817,8 @@ class StableDiffusionIPEXPipeline(DiffusionPipeline): ...@@ -817,7 +817,8 @@ class StableDiffusionIPEXPipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if output_type == "latent": if output_type == "latent":
image = latents image = latents
......
...@@ -770,7 +770,8 @@ class StableDiffusionReferencePipeline(StableDiffusionPipeline): ...@@ -770,7 +770,8 @@ class StableDiffusionReferencePipeline(StableDiffusionPipeline):
if i == len(timesteps) - 1 or ((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() progress_bar.update()
if callback is not None and i % callback_steps == 0: if callback is not None and i % callback_steps == 0:
callback(i, t, latents) step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent": if not output_type == "latent":
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment