Fix formula for noise levels in Karras scheduler and tests (#627)

fix formula for noise levels in karras scheduler and tests

Author: sgrigory

src/diffusers/schedulers/scheduling_karras_ve.py

@@ -110,7 +110,7 @@ def set_timesteps(self, num_inference_steps: int):
         self.timesteps = np.arange(0, self.num_inference_steps)[::-1].copy()
         self.schedule = [
             (
-                self.config.sigma_max
+                self.config.sigma_max**2
                 * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
             )
             for i in self.timesteps

src/diffusers/schedulers/scheduling_karras_ve_flax.py

@@ -113,7 +113,7 @@ def set_timesteps(self, state: KarrasVeSchedulerState, num_inference_steps: int)
         timesteps = jnp.arange(0, num_inference_steps)[::-1].copy()
         schedule = [
             (
-                self.config.sigma_max
+                self.config.sigma_max**2
                 * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
             )
             for i in timesteps

tests/test_pipelines.py

@@ -1104,7 +1104,7 @@ def test_karras_ve_pipeline(self):
 
         image_slice = image[0, -3:, -3:, -1]
         assert image.shape == (1, 256, 256, 3)
-        expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837])
+        expected_slice = np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586])
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
 
     @slow