From e8923c32a5a64ccc4c62699f5779433907ac37c9 Mon Sep 17 00:00:00 2001 From: Vuln3rat0r Date: Fri, 23 Oct 2026 16:03:14 +0000 Subject: [PATCH 1/2] fix(schedulers.lms): bound multistep order by available derivatives (#15108) --- .../schedulers/scheduling_lms_discrete.py | 2 +- tests/schedulers/test_scheduler_lms.py | 30 +++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index e917adda9516..9437ea0af687 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -603,7 +603,7 @@ def step( self.derivatives.pop(0) # 3. Compute linear multistep coefficients - order = min(self.step_index + 1, order) + order = min(self.step_index + 1, len(self.derivatives), order) lms_coeffs = [self.get_lms_coefficient(order, self.step_index, curr_order) for curr_order in range(order)] # 4. Compute previous sample based on the derivatives path diff --git a/tests/schedulers/test_scheduler_lms.py b/tests/schedulers/test_scheduler_lms.py index c4abca3ac973..f5f1cbacaabd 100644 --- a/tests/schedulers/test_scheduler_lms.py +++ b/tests/schedulers/test_scheduler_lms.py @@ -166,8 +166,34 @@ def test_full_loop_with_noise(self): result_sum = torch.sum(torch.abs(sample)) result_mean = torch.mean(torch.abs(sample)) - assert abs(result_sum.item() - 27663.6895) < 1e-2 - assert abs(result_mean.item() - 36.0204) < 1e-3 + assert abs(result_sum.item() - 57074.7617) < 1e-2 + assert abs(result_mean.item() - 74.3161) < 1e-3 + + def test_step_order_with_begin_index(self): + scheduler_class = self.scheduler_classes[0] + scheduler_config = self.get_scheduler_config() + scheduler = scheduler_class(**scheduler_config) + + scheduler.set_timesteps(self.num_inference_steps) + scheduler.set_begin_index(5) + + sample = self.dummy_sample_deter + model_output = self.dummy_noise_deter + + # On the first step starting at begin_index, only 1 derivative is available + # so the update must match first-order Euler step (order=1) + sigma_curr = scheduler.sigmas[5] + sigma_next = scheduler.sigmas[6] + derivative = (sample - (sample - sigma_curr * model_output)) / sigma_curr + expected_prev = sample + (sigma_next - sigma_curr) * derivative + + output = scheduler.step(model_output, scheduler.timesteps[5], sample) + assert torch.allclose(output.prev_sample, expected_prev, atol=1e-5) + assert len(scheduler.derivatives) == 1 + + # Second step accumulates a second derivative + scheduler.step(model_output, scheduler.timesteps[6], output.prev_sample) + assert len(scheduler.derivatives) == 2 def test_beta_sigmas(self): self.check_over_configs(use_beta_sigmas=True) From fe5db08e821bb5275621fa92e137958d19892a97 Mon Sep 17 00:00:00 2001 From: Vuln3rat0r Date: Thu, 5 Nov 2026 21:47:44 +0000 Subject: [PATCH 2/2] test(schedulers.lms): verify numerical updates for steps 2 and 3 on resumed run --- tests/schedulers/test_scheduler_lms.py | 42 ++++++++++++++++++-------- 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/tests/schedulers/test_scheduler_lms.py b/tests/schedulers/test_scheduler_lms.py index f5f1cbacaabd..b68472d63057 100644 --- a/tests/schedulers/test_scheduler_lms.py +++ b/tests/schedulers/test_scheduler_lms.py @@ -178,23 +178,39 @@ def test_step_order_with_begin_index(self): scheduler.set_begin_index(5) sample = self.dummy_sample_deter - model_output = self.dummy_noise_deter - - # On the first step starting at begin_index, only 1 derivative is available - # so the update must match first-order Euler step (order=1) - sigma_curr = scheduler.sigmas[5] - sigma_next = scheduler.sigmas[6] - derivative = (sample - (sample - sigma_curr * model_output)) / sigma_curr - expected_prev = sample + (sigma_next - sigma_curr) * derivative - - output = scheduler.step(model_output, scheduler.timesteps[5], sample) - assert torch.allclose(output.prev_sample, expected_prev, atol=1e-5) + model_output_1 = self.dummy_noise_deter + model_output_2 = self.dummy_noise_deter * 1.5 + 0.1 + model_output_3 = self.dummy_noise_deter * 0.8 - 0.2 + + # Step 1: Only 1 derivative is available, update must match first-order Euler step (order=1) + sigma_5 = scheduler.sigmas[5] + sigma_6 = scheduler.sigmas[6] + d1 = (sample - (sample - sigma_5 * model_output_1)) / sigma_5 + expected_prev_1 = sample + (sigma_6 - sigma_5) * d1 + + output_1 = scheduler.step(model_output_1, scheduler.timesteps[5], sample) + assert torch.allclose(output_1.prev_sample, expected_prev_1, atol=1e-5) assert len(scheduler.derivatives) == 1 - # Second step accumulates a second derivative - scheduler.step(model_output, scheduler.timesteps[6], output.prev_sample) + # Step 2: 2 derivatives available, update must match second-order Adams-Bashforth (order=2) + d2 = (output_1.prev_sample - (output_1.prev_sample - sigma_6 * model_output_2)) / sigma_6 + coeffs_2 = [scheduler.get_lms_coefficient(2, 6, c) for c in range(2)] + expected_prev_2 = output_1.prev_sample + coeffs_2[0] * d2 + coeffs_2[1] * d1 + + output_2 = scheduler.step(model_output_2, scheduler.timesteps[6], output_1.prev_sample) + assert torch.allclose(output_2.prev_sample, expected_prev_2, atol=1e-5) assert len(scheduler.derivatives) == 2 + # Step 3: 3 derivatives available, update must match third-order Adams-Bashforth (order=3) + sigma_7 = scheduler.sigmas[7] + d3 = (output_2.prev_sample - (output_2.prev_sample - sigma_7 * model_output_3)) / sigma_7 + coeffs_3 = [scheduler.get_lms_coefficient(3, 7, c) for c in range(3)] + expected_prev_3 = output_2.prev_sample + coeffs_3[0] * d3 + coeffs_3[1] * d2 + coeffs_3[2] * d1 + + output_3 = scheduler.step(model_output_3, scheduler.timesteps[7], output_2.prev_sample) + assert torch.allclose(output_3.prev_sample, expected_prev_3, atol=1e-5) + assert len(scheduler.derivatives) == 3 + def test_beta_sigmas(self): self.check_over_configs(use_beta_sigmas=True)