Skip to content

Commit

Permalink
fix(diffusers/schedulers): fix some bug in schedulers
Browse files Browse the repository at this point in the history
  • Loading branch information
Cui-yshoho committed Nov 7, 2024
1 parent 29ca665 commit 96aa40b
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 6 deletions.
1 change: 1 addition & 0 deletions mindone/diffusers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@
"DPMSolverMultistepScheduler",
"DPMSolverMultistepInverseScheduler",
"DPMSolverSinglestepScheduler",
"EDMDPMSolverMultistepScheduler",
"EDMEulerScheduler",
"EulerAncestralDiscreteScheduler",
"EulerDiscreteScheduler",
Expand Down
16 changes: 10 additions & 6 deletions mindone/diffusers/schedulers/scheduling_edm_euler.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,18 +249,22 @@ def index_for_timestep(self, timestep, schedule_timesteps=None):
if schedule_timesteps is None:
schedule_timesteps = self.timesteps

if (schedule_timesteps == timestep).sum() > 1:
pos = 1
else:
pos = 0
index_candidates_num = (schedule_timesteps == timestep).sum()

if index_candidates_num == 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)
indices = (schedule_timesteps == timestep).nonzero()
else:
if index_candidates_num > 1:
pos = 1
else:
pos = 0
step_index = int((schedule_timesteps == timestep).nonzero()[pos])

return int(indices[pos])
return step_index

# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
def _init_step_index(self, timestep):
Expand Down

0 comments on commit 96aa40b

Please sign in to comment.