Remove additional noise in final sample of DDIM inference process

finegrain-ai · Jan 18, 2024 · 17d9701 · 17d9701
1 parent a1f50f3
commit 17d9701
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/refiners/foundationals/latent_diffusion/schedulers/ddim.py b/src/refiners/foundationals/latent_diffusion/schedulers/ddim.py
@@ -52,6 +52,12 @@ def __call__(self, x: Tensor, noise: Tensor, step: int, generator: Generator | N
             ),
         )
         predicted_x = (x - sqrt(1 - current_scale_factor**2) * noise) / current_scale_factor
-        denoised_x = previous_scale_factor * predicted_x + sqrt(1 - previous_scale_factor**2) * noise
+        noise_factor = sqrt(1 - previous_scale_factor**2)
+
+        # Do not add noise at the last step to avoid visual artifacts.
+        if step == self.num_inference_steps - 1:
+            noise_factor = 0
+
+        denoised_x = previous_scale_factor * predicted_x + noise_factor * noise
 
         return denoised_x
diff --git a/tests/foundationals/latent_diffusion/test_schedulers.py b/tests/foundationals/latent_diffusion/test_schedulers.py
@@ -46,7 +46,6 @@ def test_ddim_diffusers():
         beta_schedule="scaled_linear",
         beta_start=0.00085,
         num_train_timesteps=1000,
-        set_alpha_to_one=False,
         steps_offset=1,
         clip_sample=False,
     )
@@ -103,7 +102,6 @@ def test_scheduler_remove_noise():
         beta_schedule="scaled_linear",
         beta_start=0.00085,
         num_train_timesteps=1000,
-        set_alpha_to_one=False,
         steps_offset=1,
         clip_sample=False,
     )