Skip to content

fix a bug in 2nd order schedulers when using in ensemble of experts config #5511

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Oct 25, 2023
Merged

fix a bug in 2nd order schedulers when using in ensemble of experts config #5511

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
4b1e7d2
fix
Oct 24, 2023
8ea162b
fix copies
Oct 24, 2023
5a0ac58
remove heun from tests
Oct 24, 2023
443aa6b
add back heun and fix the tests to include 2nd order
Oct 24, 2023
de85e0b
fix the other test too
Oct 24, 2023
ceb0f45
Apply suggestions from code review
patrickvonplaten Oct 25, 2023
62a9b6e
Apply suggestions from code review
patrickvonplaten Oct 25, 2023
a424a3f
Apply suggestions from code review
patrickvonplaten Oct 25, 2023
ea590b8
make style
patrickvonplaten Oct 25, 2023
2b959df
add more comments
patrickvonplaten Oct 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -867,8 +867,20 @@ def get_timesteps(self, num_inference_steps, strength, device, denoising_start=N
- (denoising_start * self.scheduler.config.num_train_timesteps)
)
)
timesteps = list(filter(lambda ts: ts < discrete_timestep_cutoff, timesteps))
return torch.tensor(timesteps), len(timesteps)

num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
if self.scheduler.order == 2:
# if the scheduler is a 2nd order scheduler we ALWAYS have to do +1
# because `num_inference_steps` will always be even given that every timestep
# (except the highest one) is duplicated. If `num_inference_steps` is even it would
# mean that we cut the timesteps in the middle of the denoising step
# (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
# we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
num_inference_steps = num_inference_steps + 1
Copy link

@nhnt11 nhnt11 Oct 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_inference_steps will always be even. Consider the case where the scheduler has been set up by calling set_timesteps(25) and then supplied denoising_start = 0.6. The following code prints an even number:

sched = SCHEDULERS["DPM2"]
sched.set_timesteps(25)
denoising_start = 0.6
discrete_timestep_cutoff = int(
    round(
        sched.config.num_train_timesteps
        - (denoising_start * sched.config.num_train_timesteps)
    )
)
num_inference_steps = (sched.timesteps < discrete_timestep_cutoff).sum().item()
print(num_inference_steps)

This prints 20. For denoising_start=0.5, it prints 25. So num_inference_steps may be either odd or even depending on num_inference_steps and denoising_start passed into the pipeline call.

I think we should add 1 to it only if it's even, and not if it's odd. No?

@yiyixuxu @patrickvonplaten

Copy link
Contributor

@patrickvonplaten patrickvonplaten Oct 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah DPM2 can actually be different indeed if it interpolates, (Heun can't)

Nice catch! I'll open a PR to fix it

Copy link

@nhnt11 nhnt11 Oct 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the result I get with hnf = 0.5 and steps = 25 with DPM2 with the code from this PR:
image

Copy link
Contributor

@patrickvonplaten patrickvonplaten Oct 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#5526 (comment)

Copy link

@nhnt11 nhnt11 Oct 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @patrickvonplaten, all testcases I threw at it now produce good images.

I am just curious why we don't need to do any logic in the base step (where denoising_end is supplied to the base SDXL pipeline without img2img) since I would expect that we would want to ensure that the final timestep there is a second order timestep, and as far as I can tell it still has the odd/even issue when splitting. Maybe I am wrong!

In any case, latest main now produces good images for all my testcases like I said. Thank you! 🙏


# because t_n+1 >= t_n, we slice the timesteps starting from the end
timesteps = timesteps[-num_inference_steps:]
return timesteps, num_inference_steps

return timesteps, num_inference_steps - t_start

Expand Down
16 changes: 14 additions & 2 deletions src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -515,8 +515,20 @@ def get_timesteps(self, num_inference_steps, strength, device, denoising_start=N
- (denoising_start * self.scheduler.config.num_train_timesteps)
)
)
timesteps = list(filter(lambda ts: ts < discrete_timestep_cutoff, timesteps))
return torch.tensor(timesteps), len(timesteps)

num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
if self.scheduler.order == 2:
# if the scheduler is a 2nd order scheduler we ALWAYS have to do +1
# because `num_inference_steps` will always be even given that every timestep
# (except the highest one) is duplicated. If `num_inference_steps` is even it would
# mean that we cut the timesteps in the middle of the denoising step
# (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
# we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
num_inference_steps = num_inference_steps + 1

# because t_n+1 >= t_n, we slice the timesteps starting from the end
timesteps = timesteps[-num_inference_steps:]
return timesteps, num_inference_steps

return timesteps, num_inference_steps - t_start

Expand Down
16 changes: 14 additions & 2 deletions src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,8 +799,20 @@ def get_timesteps(self, num_inference_steps, strength, device, denoising_start=N
- (denoising_start * self.scheduler.config.num_train_timesteps)
)
)
timesteps = list(filter(lambda ts: ts < discrete_timestep_cutoff, timesteps))
return torch.tensor(timesteps), len(timesteps)

num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
if self.scheduler.order == 2:
# if the scheduler is a 2nd order scheduler we ALWAYS have to do +1
# because `num_inference_steps` will always be even given that every timestep
# (except the highest one) is duplicated. If `num_inference_steps` is even it would
# mean that we cut the timesteps in the middle of the denoising step
# (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
# we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
num_inference_steps = num_inference_steps + 1

# because t_n+1 >= t_n, we slice the timesteps starting from the end
timesteps = timesteps[-num_inference_steps:]
return timesteps, num_inference_steps

return timesteps, num_inference_steps - t_start

Expand Down
26 changes: 18 additions & 8 deletions tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,8 +324,13 @@ class scheduler_cls(scheduler_cls_orig):
pipe_1.scheduler.set_timesteps(num_steps)
expected_steps = pipe_1.scheduler.timesteps.tolist()

expected_steps_1 = list(filter(lambda ts: ts >= split, expected_tss))
expected_steps_2 = list(filter(lambda ts: ts < split, expected_tss))
if pipe_1.scheduler.order == 2:
expected_steps_1 = list(filter(lambda ts: ts >= split, expected_tss))
expected_steps_2 = expected_steps_1[-1:] + list(filter(lambda ts: ts < split, expected_tss))
expected_steps = expected_steps_1 + expected_steps_2
else:
expected_steps_1 = list(filter(lambda ts: ts >= split, expected_tss))
expected_steps_2 = list(filter(lambda ts: ts < split, expected_tss))

# now we monkey patch step `done_steps`
# list into the step function for testing
Expand Down Expand Up @@ -607,13 +612,18 @@ class scheduler_cls(scheduler_cls_orig):

split_1_ts = num_train_timesteps - int(round(num_train_timesteps * split_1))
split_2_ts = num_train_timesteps - int(round(num_train_timesteps * split_2))
expected_steps_1 = expected_steps[:split_1_ts]
expected_steps_2 = expected_steps[split_1_ts:split_2_ts]
expected_steps_3 = expected_steps[split_2_ts:]

expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps))
expected_steps_3 = list(filter(lambda ts: ts < split_2_ts, expected_steps))
if pipe_1.scheduler.order == 2:
expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = expected_steps_1[-1:] + list(
filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps)
)
expected_steps_3 = expected_steps_2[-1:] + list(filter(lambda ts: ts < split_2_ts, expected_steps))
expected_steps = expected_steps_1 + expected_steps_2 + expected_steps_3
else:
expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps))
expected_steps_3 = list(filter(lambda ts: ts < split_2_ts, expected_steps))

# now we monkey patch step `done_steps`
# list into the step function for testing
Expand Down
28 changes: 18 additions & 10 deletions tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,11 +318,14 @@ class scheduler_cls(scheduler_cls_orig):
expected_steps = pipe_1.scheduler.timesteps.tolist()

split_ts = num_train_timesteps - int(round(num_train_timesteps * split))
expected_steps_1 = expected_steps[:split_ts]
expected_steps_2 = expected_steps[split_ts:]

expected_steps_1 = list(filter(lambda ts: ts >= split_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts < split_ts, expected_steps))
if pipe_1.scheduler.order == 2:
expected_steps_1 = list(filter(lambda ts: ts >= split_ts, expected_steps))
expected_steps_2 = expected_steps_1[-1:] + list(filter(lambda ts: ts < split_ts, expected_steps))
expected_steps = expected_steps_1 + expected_steps_2
else:
expected_steps_1 = list(filter(lambda ts: ts >= split_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts < split_ts, expected_steps))

# now we monkey patch step `done_steps`
# list into the step function for testing
Expand Down Expand Up @@ -389,13 +392,18 @@ class scheduler_cls(scheduler_cls_orig):

split_1_ts = num_train_timesteps - int(round(num_train_timesteps * split_1))
split_2_ts = num_train_timesteps - int(round(num_train_timesteps * split_2))
expected_steps_1 = expected_steps[:split_1_ts]
expected_steps_2 = expected_steps[split_1_ts:split_2_ts]
expected_steps_3 = expected_steps[split_2_ts:]

expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps))
expected_steps_3 = list(filter(lambda ts: ts < split_2_ts, expected_steps))
if pipe_1.scheduler.order == 2:
expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = expected_steps_1[-1:] + list(
filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps)
)
expected_steps_3 = expected_steps_2[-1:] + list(filter(lambda ts: ts < split_2_ts, expected_steps))
expected_steps = expected_steps_1 + expected_steps_2 + expected_steps_3
else:
expected_steps_1 = list(filter(lambda ts: ts >= split_1_ts, expected_steps))
expected_steps_2 = list(filter(lambda ts: ts >= split_2_ts and ts < split_1_ts, expected_steps))
expected_steps_3 = list(filter(lambda ts: ts < split_2_ts, expected_steps))

# now we monkey patch step `done_steps`
# list into the step function for testing
Expand Down