From 615baa600010ae44f479d983a77d0cd1037ff24c Mon Sep 17 00:00:00 2001 From: hugojarkoff Date: Thu, 18 Jan 2024 18:31:00 +0100 Subject: [PATCH] Remove additional noise in final sample of DDIM inference process --- .../foundationals/latent_diffusion/schedulers/ddim.py | 8 +++++++- tests/foundationals/latent_diffusion/test_schedulers.py | 2 -- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/refiners/foundationals/latent_diffusion/schedulers/ddim.py b/src/refiners/foundationals/latent_diffusion/schedulers/ddim.py index f6f2cba15..3ff1d8382 100644 --- 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 index 088b578a8..a6a524ef1 100644 --- 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, )