Skip to content

Fix zero-step NaN and scalar sync in EDM, Cosine, and Inverse DPMSolverMultistepScheduler #15080

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

Open
deep1blue wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix zero-step NaN and scalar sync in EDM, Cosine, and Inverse DPMSolverMultistepScheduler #15080

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 10 additions & 6 deletions src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,15 +280,15 @@ def set_timesteps(self, num_inference_steps: int | None = None, device: str | to
self.timesteps = self.precondition_noise(sigmas)

if self.config.final_sigmas_type == "sigma_min":
sigma_last = self.config.sigma_min
sigma_last = sigmas[-1:]
elif self.config.final_sigmas_type == "zero":
sigma_last = 0
sigma_last = torch.zeros(1, dtype=torch.float32, device=device)
else:
raise ValueError(
f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
)

self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)])
self.sigmas = torch.cat([sigmas, sigma_last])

self.model_outputs = [
None,
Expand Down Expand Up @@ -643,6 +643,7 @@ def step(
self.config.euler_at_final
or (self.config.lower_order_final and len(self.timesteps) < 15)
or self.config.final_sigmas_type == "zero"
or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index])
)
lower_order_second = (
(self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
Expand All @@ -665,9 +666,12 @@ def step(
sigma_max=self.config.sigma_max,
seed=seed,
)
noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to(
model_output.device
)
if self.sigmas[self.step_index + 1] == self.sigmas[self.step_index]:
noise = torch.zeros_like(model_output)
else:
noise = self.noise_sampler(self.sigmas[self.step_index], self.sigmas[self.step_index + 1]).to(
model_output.device
)

if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise)
Expand Down
4 changes: 3 additions & 1 deletion src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1058,7 +1058,9 @@ def step(

# Improve numerical stability for small number of steps
lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
self.config.euler_at_final or (self.config.lower_order_final and len(self.timesteps) < 15)
self.config.euler_at_final
or (self.config.lower_order_final and len(self.timesteps) < 15)
or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index])
)
lower_order_second = (
(self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
Expand Down
7 changes: 4 additions & 3 deletions src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,15 +303,15 @@ def set_timesteps(
self.timesteps = self.precondition_noise(sigmas)

if self.config.final_sigmas_type == "sigma_min":
sigma_last = self.config.sigma_min
sigma_last = sigmas[-1:]
elif self.config.final_sigmas_type == "zero":
sigma_last = 0
sigma_last = torch.zeros(1, dtype=torch.float32, device=device)
else:
raise ValueError(
f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
)

self.sigmas = torch.cat([sigmas, torch.tensor([sigma_last], dtype=torch.float32, device=device)])
self.sigmas = torch.cat([sigmas, sigma_last])

self.model_outputs = [
None,
Expand Down Expand Up @@ -766,6 +766,7 @@ def step(
self.config.euler_at_final
or (self.config.lower_order_final and len(self.timesteps) < 15)
or self.config.final_sigmas_type == "zero"
or bool(self.sigmas[self.step_index + 1] == self.sigmas[self.step_index])
)
lower_order_second = (
(self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
Expand Down
20 changes: 20 additions & 0 deletions tests/schedulers/test_scheduler_dpm_multi_inverse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,3 +271,23 @@ def test_beta_sigmas(self):

def test_exponential_sigmas(self):
self.check_over_configs(use_exponential_sigmas=True)

def test_zero_width_final_step(self):
scheduler_class = self.scheduler_classes[0]
model = self.dummy_model()
for kwargs in [
{"solver_order": 3, "use_karras_sigmas": True},
{"solver_order": 3, "use_exponential_sigmas": True},
{"solver_order": 3, "use_beta_sigmas": True},
{"solver_order": 3, "timestep_spacing": "trailing"},
{"solver_order": 2, "solver_type": "heun", "use_karras_sigmas": True},
]:
scheduler_config = self.get_scheduler_config(**kwargs)
scheduler = scheduler_class(**scheduler_config)
scheduler.set_timesteps(20)
sample = self.dummy_sample_deter
generator = torch.manual_seed(0)
for t in scheduler.timesteps:
residual = model(sample, t)
sample = scheduler.step(residual, t, sample, generator=generator).prev_sample
assert not torch.isnan(sample).any()
53 changes: 51 additions & 2 deletions tests/schedulers/test_scheduler_edm_dpmsolver_multistep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,6 @@ def test_prediction_type(self):
for prediction_type in ["epsilon", "v_prediction"]:
self.check_over_configs(prediction_type=prediction_type)

# TODO (patil-suraj): Fix this test
@unittest.skip("Skip for now, as it failing currently but works with the actual model")
def test_solver_order_and_type(self):
for algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
for solver_type in ["midpoint", "heun"]:
Expand Down Expand Up @@ -260,3 +258,54 @@ def test_duplicated_timesteps(self, **config):
@unittest.skip("Test not supported.")
def test_trained_betas(self):
pass

def test_zero_width_final_step(self):
scheduler_class = self.scheduler_classes[0]
model = self.dummy_model()
for kwargs in [
{"solver_order": 3, "sigma_schedule": "karras", "final_sigmas_type": "sigma_min"},
{"solver_order": 3, "sigma_schedule": "exponential", "final_sigmas_type": "sigma_min"},
{"solver_order": 2, "solver_type": "heun", "sigma_schedule": "karras", "final_sigmas_type": "sigma_min"},
{
"solver_order": 2,
"solver_type": "heun",
"algorithm_type": "sde-dpmsolver++",
"sigma_min": 0.01,
"final_sigmas_type": "sigma_min",
},
{
"solver_order": 1,
"algorithm_type": "sde-dpmsolver++",
"sigma_min": 0.01,
"final_sigmas_type": "sigma_min",
},
]:
scheduler_config = self.get_scheduler_config(**kwargs)
scheduler = scheduler_class(**scheduler_config)
scheduler.set_timesteps(20)
sample = self.dummy_sample_deter
generator = torch.manual_seed(0)
for t in scheduler.timesteps:
residual = model(sample, t)
sample = scheduler.step(residual, t, sample, generator=generator).prev_sample
assert not torch.isnan(sample).any()

from diffusers import CosineDPMSolverMultistepScheduler

for solver_order in [1, 2]:
for solver_type in ["midpoint", "heun"]:
cosine_scheduler = CosineDPMSolverMultistepScheduler(
solver_order=solver_order,
solver_type=solver_type,
final_sigmas_type="sigma_min",
)
cosine_scheduler.set_timesteps(10)
sample = self.dummy_sample_deter
# BrownianTreeNoiseSampler divides by sqrt(|t1 - t0|) and yields NaN when sigma == sigma_next
cosine_scheduler.noise_sampler = lambda s, s_next: (
torch.full_like(sample, float("nan")) if s == s_next else torch.zeros_like(sample)
)
for t in cosine_scheduler.timesteps:
residual = model(sample, t)
sample = cosine_scheduler.step(residual, t, sample).prev_sample
assert not torch.isnan(sample).any()
Loading