diffusers#4003 Correct the method being used for skipping steps, and the count of steps

Author: bghira

src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -533,13 +533,20 @@ def get_timesteps(self, num_inference_steps, begin_inference_step, strength, dev
         # get the original timestep using init_timestep
         if begin_inference_step is None:
             init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+            t_start = max(num_inference_steps - init_timestep, 0)
+            new_num_inference_steps = num_inference_steps - t_start
+            target_count = t_start * self.scheduler.order
         else:
-            init_timestep = begin_inference_step - 1
+            init_timestep = begin_inference_step
+            t_start = max(num_inference_steps - init_timestep, 0)
+            new_num_inference_steps = num_inference_steps - (num_inference_steps - t_start)
+            target_count = num_inference_steps
+        logger.info(f'Calculating t_start via max({num_inference_steps} - {init_timestep}, 0)')
+        logger.info(f't_start is {t_start}')
+        logger.info(f'We will have {target_count} timesteps from {len(self.scheduler.timesteps)} total')
+        timesteps = self.scheduler.timesteps[target_count :]
 
-        t_start = max(num_inference_steps - init_timestep, 0)
-        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
-
-        return timesteps, num_inference_steps - t_start
+        return timesteps, new_num_inference_steps
 
     def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True):
         if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
@@ -924,6 +931,8 @@ def __call__(
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
+                if begin_inference_step is not None and i < begin_inference_step:
+                    logger.info(f'Skipping step {i} because we are waiting for {begin_inference_step}')
                 # expand the latents if we are doing classifier free guidance
                 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents