From 466d601c8c1e0aae063d9bea8c17d9356d31ad90 Mon Sep 17 00:00:00 2001 From: Muhammad Mirza <114640561+mirzasnowy@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:49:23 +0700 Subject: [PATCH 1/2] fix(schedulers): skip noise at terminal DDPM step when no predecessor exists With timestep_spacing='trailing', the final selected timestep is nonzero, so DDPMScheduler.step() and DDPMParallelScheduler.step() added fresh noise at the terminal transition even though the mean is already the predicted clean sample. Gate noise addition on predecessor existence (prev_t >= 0) instead of on t > 0, matching how the same methods already handle the terminal case for the alpha coefficients. Fixes #15139 --- src/diffusers/schedulers/scheduling_ddpm.py | 2 +- .../schedulers/scheduling_ddpm_parallel.py | 2 +- tests/schedulers/test_scheduler_ddpm.py | 28 +++++++++++++++++++ .../test_scheduler_ddpm_parallel.py | 28 +++++++++++++++++++ 4 files changed, 58 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_ddpm.py b/src/diffusers/schedulers/scheduling_ddpm.py index 972c46c6e930..d8f4f2c6f569 100644 --- a/src/diffusers/schedulers/scheduling_ddpm.py +++ b/src/diffusers/schedulers/scheduling_ddpm.py @@ -540,7 +540,7 @@ def step( # 6. Add noise variance = 0 - if t > 0: + if prev_t >= 0: device = model_output.device variance_noise = randn_tensor( model_output.shape, diff --git a/src/diffusers/schedulers/scheduling_ddpm_parallel.py b/src/diffusers/schedulers/scheduling_ddpm_parallel.py index 725958b01a44..f57a1389e414 100644 --- a/src/diffusers/schedulers/scheduling_ddpm_parallel.py +++ b/src/diffusers/schedulers/scheduling_ddpm_parallel.py @@ -555,7 +555,7 @@ def step( # 6. Add noise variance = 0 - if t > 0: + if prev_t >= 0: device = model_output.device variance_noise = randn_tensor( model_output.shape, diff --git a/tests/schedulers/test_scheduler_ddpm.py b/tests/schedulers/test_scheduler_ddpm.py index 056b5d83350e..bff1d816c4b0 100644 --- a/tests/schedulers/test_scheduler_ddpm.py +++ b/tests/schedulers/test_scheduler_ddpm.py @@ -68,6 +68,34 @@ def test_variance(self): assert torch.sum(torch.abs(scheduler._get_variance(487) - 0.00979)) < 1e-5 assert torch.sum(torch.abs(scheduler._get_variance(999) - 0.02)) < 1e-5 + def test_terminal_step_adds_no_noise(self): + # Regression test for https://github.com/huggingface/diffusers/issues/15139: + # with trailing timestep spacing the final selected timestep is nonzero, + # so `step()` must not add fresh noise at the terminal transition. + scheduler_class = self.scheduler_classes[0] + sample = torch.zeros(1, 1, 2, 2) + pred_original_sample = torch.full_like(sample, 0.25) + + for variance_type in ["fixed_small", "fixed_large"]: + scheduler_config = self.get_scheduler_config( + variance_type=variance_type, + timestep_spacing="trailing", + prediction_type="sample", + clip_sample=False, + ) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(50) + t = scheduler.timesteps[-1] + assert scheduler.previous_timestep(t) == -1 + + output = scheduler.step( + pred_original_sample, + t, + sample, + generator=torch.Generator().manual_seed(0), + ) + assert torch.equal(output.prev_sample, output.pred_original_sample) + def test_rescale_betas_zero_snr(self): for rescale_betas_zero_snr in [True, False]: self.check_over_configs(rescale_betas_zero_snr=rescale_betas_zero_snr) diff --git a/tests/schedulers/test_scheduler_ddpm_parallel.py b/tests/schedulers/test_scheduler_ddpm_parallel.py index 377067071c25..da3ba387a4d7 100644 --- a/tests/schedulers/test_scheduler_ddpm_parallel.py +++ b/tests/schedulers/test_scheduler_ddpm_parallel.py @@ -82,6 +82,34 @@ def test_variance(self): assert torch.sum(torch.abs(scheduler._get_variance(487) - 0.00979)) < 1e-5 assert torch.sum(torch.abs(scheduler._get_variance(999) - 0.02)) < 1e-5 + def test_terminal_step_adds_no_noise(self): + # Regression test for https://github.com/huggingface/diffusers/issues/15139: + # with trailing timestep spacing the final selected timestep is nonzero, + # so `step()` must not add fresh noise at the terminal transition. + scheduler_class = self.scheduler_classes[0] + sample = torch.zeros(1, 1, 2, 2) + pred_original_sample = torch.full_like(sample, 0.25) + + for variance_type in ["fixed_small", "fixed_large"]: + scheduler_config = self.get_scheduler_config( + variance_type=variance_type, + timestep_spacing="trailing", + prediction_type="sample", + clip_sample=False, + ) + scheduler = scheduler_class(**scheduler_config) + scheduler.set_timesteps(50) + t = scheduler.timesteps[-1] + assert scheduler.previous_timestep(t) == -1 + + output = scheduler.step( + pred_original_sample, + t, + sample, + generator=torch.Generator().manual_seed(0), + ) + assert torch.equal(output.prev_sample, output.pred_original_sample) + def test_rescale_betas_zero_snr(self): for rescale_betas_zero_snr in [True, False]: self.check_over_configs(rescale_betas_zero_snr=rescale_betas_zero_snr) From 5c5fcb2f595d63fcbf4aba2ee6fffb1821606693 Mon Sep 17 00:00:00 2001 From: Muhammad Mirza <114640561+mirzasnowy@users.noreply.github.com> Date: Wed, 4 Nov 2026 03:16:26 +0700 Subject: [PATCH 2/2] test(schedulers): assert generator state unchanged at terminal DDPM step Addresses ZihaoW123's review suggestion on #15142: output equality alone can miss an unnecessary draw when variance is zero or shrinks, so the terminal-step test now also asserts the generator state is unchanged. Scheduler code is untouched; only the two DDPM test files change. --- tests/schedulers/test_scheduler_ddpm.py | 8 +++++++- tests/schedulers/test_scheduler_ddpm_parallel.py | 8 +++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/schedulers/test_scheduler_ddpm.py b/tests/schedulers/test_scheduler_ddpm.py index bff1d816c4b0..83df337b9147 100644 --- a/tests/schedulers/test_scheduler_ddpm.py +++ b/tests/schedulers/test_scheduler_ddpm.py @@ -88,13 +88,19 @@ def test_terminal_step_adds_no_noise(self): t = scheduler.timesteps[-1] assert scheduler.previous_timestep(t) == -1 + generator = torch.Generator().manual_seed(0) + initial_state = generator.get_state() output = scheduler.step( pred_original_sample, t, sample, - generator=torch.Generator().manual_seed(0), + generator=generator, ) assert torch.equal(output.prev_sample, output.pred_original_sample) + # The terminal step must not draw at all: output equality alone + # would miss a draw whose variance is zero or rounds away, so + # also require the supplied generator to be untouched. + assert torch.equal(generator.get_state(), initial_state) def test_rescale_betas_zero_snr(self): for rescale_betas_zero_snr in [True, False]: diff --git a/tests/schedulers/test_scheduler_ddpm_parallel.py b/tests/schedulers/test_scheduler_ddpm_parallel.py index da3ba387a4d7..5de6441fb40b 100644 --- a/tests/schedulers/test_scheduler_ddpm_parallel.py +++ b/tests/schedulers/test_scheduler_ddpm_parallel.py @@ -102,13 +102,19 @@ def test_terminal_step_adds_no_noise(self): t = scheduler.timesteps[-1] assert scheduler.previous_timestep(t) == -1 + generator = torch.Generator().manual_seed(0) + initial_state = generator.get_state() output = scheduler.step( pred_original_sample, t, sample, - generator=torch.Generator().manual_seed(0), + generator=generator, ) assert torch.equal(output.prev_sample, output.pred_original_sample) + # The terminal step must not draw at all: output equality alone + # would miss a draw whose variance is zero or rounds away, so + # also require the supplied generator to be untouched. + assert torch.equal(generator.get_state(), initial_state) def test_rescale_betas_zero_snr(self): for rescale_betas_zero_snr in [True, False]: