From 4633e7ed05adb91b556e200913c0c093040f94a7 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 14 Jan 2024 21:59:48 -0600 Subject: [PATCH] fix last step errors --- api/onnx_web/diffusers/pipelines/panorama.py | 27 +++++++++++++++++++ .../diffusers/pipelines/panorama_xl.py | 17 ++++++++++-- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/api/onnx_web/diffusers/pipelines/panorama.py b/api/onnx_web/diffusers/pipelines/panorama.py index cd9a32d0..c5db3cf3 100644 --- a/api/onnx_web/diffusers/pipelines/panorama.py +++ b/api/onnx_web/diffusers/pipelines/panorama.py @@ -578,6 +578,8 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): last = i == (len(self.scheduler.timesteps) - 1) + next_step_index = None + count.fill(0) value.fill(0) @@ -633,6 +635,7 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): self.scheduler._step_index, prev_step_index, ) + next_step_index = self.scheduler._step_index self.scheduler._step_index = prev_step_index value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised @@ -721,6 +724,7 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): self.scheduler._step_index, prev_step_index, ) + next_step_index = self.scheduler._step_index self.scheduler._step_index = prev_step_index if feather[0] > 0.0: @@ -751,6 +755,16 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): latents = np.where(count > 0, value / count, value) latents = repair_nan(latents) + # update the scheduler's internal timestep + if not last and next_step_index is not None: + logger.debug( + "updating scheduler internal step index from %s to %s", + self.scheduler._step_index, + next_step_index, + ) + self.scheduler._step_index = next_step_index + next_step_index = None + # call the callback, if provided if callback is not None and i % callback_steps == 0: callback(i, t, latents) @@ -1021,6 +1035,9 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): ) for i, t in enumerate(self.progress_bar(timesteps)): + last = i == (len(timesteps) - 1) + next_step_index = None + count.fill(0) value.fill(0) @@ -1076,6 +1093,7 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): self.scheduler._step_index, prev_step_index, ) + next_step_index = self.scheduler._step_index self.scheduler._step_index = prev_step_index value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised @@ -1084,6 +1102,15 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113 latents = np.where(count > 0, value / count, value) + # update the scheduler's internal timestep + if not last and next_step_index is not None: + logger.debug( + "updating scheduler internal step index from %s to %s", + self.scheduler._step_index, + next_step_index, + ) + self.scheduler._step_index = next_step_index + # call the callback, if provided if callback is not None and i % callback_steps == 0: callback(i, t, latents) diff --git a/api/onnx_web/diffusers/pipelines/panorama_xl.py b/api/onnx_web/diffusers/pipelines/panorama_xl.py index fc41fb41..e1302426 100644 --- a/api/onnx_web/diffusers/pipelines/panorama_xl.py +++ b/api/onnx_web/diffusers/pipelines/panorama_xl.py @@ -607,9 +607,9 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix latents = repair_nan(latents) # update the scheduler's internal timestep, if set - if next_step_index is not None: + if not last and next_step_index is not None: logger.debug( - "forwarding scheduler internal step index from %s to %s", + "updating scheduler internal step index from %s to %s", self.scheduler._step_index, next_step_index, ) @@ -876,6 +876,9 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order for i, t in enumerate(self.progress_bar(timesteps)): + last = i == (len(timesteps) - 1) + next_step_index = None + count.fill(0) value.fill(0) @@ -940,6 +943,7 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix self.scheduler._step_index, prev_step_index, ) + next_step_index = self.scheduler._step_index self.scheduler._step_index = prev_step_index value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised @@ -948,6 +952,15 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113 latents = np.where(count > 0, value / count, value) + # update the scheduler's internal timestep, if set + if not last and next_step_index is not None: + logger.debug( + "updating scheduler internal step index from %s to %s", + self.scheduler._step_index, + next_step_index, + ) + self.scheduler._step_index = next_step_index + # call the callback, if provided if i == len(timesteps) - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0