diff --git a/src/diffusers/schedulers/scheduling_ddim_inverse.py b/src/diffusers/schedulers/scheduling_ddim_inverse.py index 2f1369160360..00fb526050d4 100644 --- a/src/diffusers/schedulers/scheduling_ddim_inverse.py +++ b/src/diffusers/schedulers/scheduling_ddim_inverse.py @@ -357,7 +357,14 @@ def step( pred_epsilon = model_output elif self.config.prediction_type == "sample": pred_original_sample = model_output - pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + if beta_prod_t > 0: + pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + else: + # At a zero noise level (`alpha_prod_t == 1`, e.g. the first inverse step when + # `set_alpha_to_one=True`) the sample carries no noise, so `pred_epsilon` is not identifiable from + # `sample` and the x0 prediction. Fall back to zero noise instead of dividing by zero, which would + # produce `inf` values. + pred_epsilon = torch.zeros_like(sample) elif self.config.prediction_type == "v_prediction": pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample diff --git a/tests/schedulers/test_scheduler_ddim_inverse.py b/tests/schedulers/test_scheduler_ddim_inverse.py index 81d53f1b4778..43fb18db9cca 100644 --- a/tests/schedulers/test_scheduler_ddim_inverse.py +++ b/tests/schedulers/test_scheduler_ddim_inverse.py @@ -136,3 +136,22 @@ def test_full_loop_with_no_set_alpha_to_one(self): assert abs(result_sum.item() - 542.6722) < 1e-2 assert abs(result_mean.item() - 0.7066) < 1e-3 + + def test_sample_prediction_first_step_no_inf(self): + # Regression test for #10920: with `prediction_type="sample"` and `set_alpha_to_one=True`, the first + # inverse step has `alpha_prod_t == 1` (`beta_prod_t == 0`), and the `pred_epsilon` division by + # `beta_prod_t ** 0.5` used to produce `inf` values in `prev_sample`. + torch.manual_seed(0) + scheduler = DDIMInverseScheduler(num_train_timesteps=1000, prediction_type="sample", set_alpha_to_one=True) + scheduler.set_timesteps(num_inference_steps=50) + + with torch.no_grad(): + model_output = torch.randn((1, 1, 2, 2, 2)) + sample = torch.randn((1, 1, 2, 2, 2)) + prev_sample = scheduler.step(model_output, 0, sample).prev_sample + + assert torch.isfinite(prev_sample).all() + + # And the full loop stays finite for sample prediction as well. + sample = self.full_loop(prediction_type="sample", set_alpha_to_one=True) + assert torch.isfinite(sample).all()