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):
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)
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
# 8. Post-processing
image = self.decode_latents(latents)
......
......@@ -434,7 +434,8 @@ class ImageToImageInpaintingPipeline(DiffusionPipeline):
# call the callback, if provided
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
image = self.vae.decode(latents).sample
......
......@@ -372,7 +372,8 @@ class StableDiffusionWalkPipeline(DiffusionPipeline):
# call the callback, if provided
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
image = self.vae.decode(latents).sample
......
......@@ -1088,7 +1088,8 @@ class StableDiffusionLongPromptWeightingPipeline(
progress_bar.update()
if i % callback_steps == 0:
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():
return None
......
......@@ -846,7 +846,8 @@ class OnnxStableDiffusionLongPromptWeightingPipeline(OnnxStableDiffusionPipeline
# call the callback, if provided
if i % callback_steps == 0:
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():
return None
......
......@@ -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):
progress_bar.update()
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":
# make sure the VAE is in float32 mode, as it overflows in float16
......
......@@ -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):
progress_bar.update()
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":
scaled = latents / self.vae.config.scaling_factor
......
......@@ -407,7 +407,8 @@ class MultilingualStableDiffusion(DiffusionPipeline):
# call the callback, if provided
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
image = self.vae.decode(latents).sample
......
......@@ -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):
progress_bar.update()
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
if not output_type == "latent":
......
......@@ -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):
progress_bar.update()
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
has_nsfw_concept = None
......
......@@ -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):
progress_bar.update()
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":
_latents = latents.cpu().detach().numpy() / 0.18215
......
......@@ -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):
progress_bar.update()
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":
_latents = latents.cpu().detach().numpy() / 0.18215
......
......@@ -337,7 +337,8 @@ class SeedResizeStableDiffusionPipeline(DiffusionPipeline):
# call the callback, if provided
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
image = self.vae.decode(latents).sample
......
......@@ -242,7 +242,8 @@ class SpeechToImagePipeline(DiffusionPipeline):
# call the callback, if provided
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
image = self.vae.decode(latents).sample
......
......@@ -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):
progress_bar.update()
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
# manually for max memory savings
......
......@@ -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):
progress_bar.update()
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
# manually for max memory savings
......
......@@ -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):
progress_bar.update()
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
# manually for max memory savings
......
......@@ -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):
progress_bar.update()
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
# manually for max memory savings
......
......@@ -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):
progress_bar.update()
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":
image = latents
......
......@@ -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):
progress_bar.update()
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":
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