[don't merge]Img2Img: timestep mismatch in scheduler with duplicated first timesteps

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -551,6 +551,7 @@ def get_timesteps(self, num_inference_steps, strength, device):
 
  t_start = max(num_inference_steps - init_timestep, 0)
  timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+ self.scheduler._step_index_init = t_start * self.scheduler.order
 
  return timesteps, num_inference_steps - t_start
 

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -215,13 +215,21 @@ def __init__(
  self.model_outputs = [None] * solver_order
  self.lower_order_nums = 0
  self._step_index = None
+ self._step_index_init = None
 
  @property
  def step_index(self):
  """
  The index counter for current timestep. It will increae 1 after each scheduler step.
  """
  return self._step_index
+
+ @property
+ def step_index_init(self):
+ """
+ the first step_index for denoising loop.
+ """
+ return self._step_index_init
 
  def set_timesteps(self, num_inference_steps: int = None, device: Union[str, torch.device] = None):
  """
@@ -760,23 +768,28 @@ def multistep_dpm_solver_third_order_update(
  return x_t
 
  def _init_step_index(self, timestep):
- if isinstance(timestep, torch.Tensor):
- timestep = timestep.to(self.timesteps.device)
-
- index_candidates = (self.timesteps == timestep).nonzero()
-
- if len(index_candidates) == 0:
- step_index = len(self.timesteps) - 1
- # The sigma index that is taken for the **very** first `step`
- # is always the second index (or the last index if there is only 1)
- # This way we can ensure we don't accidentally skip a sigma in
- # case we start in the middle of the denoising schedule (e.g. for image-to-image)
- elif len(index_candidates) > 1:
- step_index = index_candidates[1].item()
- else:
- step_index = index_candidates[0].item()
+
+ if self.step_index_init is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+
+ index_candidates = (self.timesteps == timestep).nonzero()
+
+ if len(index_candidates) == 0:
+ step_index = len(self.timesteps) - 1
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ elif len(index_candidates) > 1:
+ step_index = index_candidates[1].item()
+ else:
+ step_index = index_candidates[0].item()
 
- self._step_index = step_index
+ self._step_index_init = step_index
+ self._step_index = step_index
+ else:
+ self._step_index = self.step_index_init
 
  def step(
  self,
@@ -884,8 +897,10 @@ def add_noise(
  else:
  schedule_timesteps = self.timesteps.to(original_samples.device)
  timesteps = timesteps.to(original_samples.device)
-
- step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+ if self.step_index_init is None:
+ step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+ else:
+ step_indices = [self.step_index_init] * timesteps.shape[0]
 
  sigma = sigmas[step_indices].flatten()
  while len(sigma.shape) < len(original_samples.shape):