Fix DDIMInverseScheduler inf at first inverse step with prediction_type=sample

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

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()