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..b68472d63057 100644 --- a/tests/schedulers/test_scheduler_lms.py +++ b/tests/schedulers/test_scheduler_lms.py @@ -166,8 +166,50 @@ 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_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 + + # 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)