From 2a8abfaf6446b99f163ccfe58d7468e4a8174ee3 Mon Sep 17 00:00:00 2001 From: Haifeng Jin <5476582+haifeng-jin@users.noreply.github.com> Date: Wed, 17 Jan 2024 19:29:53 +0000 Subject: [PATCH] bug fix --- keras_cv/models/stable_diffusion/stable_diffusion.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/keras_cv/models/stable_diffusion/stable_diffusion.py b/keras_cv/models/stable_diffusion/stable_diffusion.py index 299f44d3d0..a68923dc78 100644 --- a/keras_cv/models/stable_diffusion/stable_diffusion.py +++ b/keras_cv/models/stable_diffusion/stable_diffusion.py @@ -209,7 +209,10 @@ def generate_image( latent = self._get_initial_diffusion_noise(batch_size, seed) # Iterative reverse diffusion stage - timesteps = np.arange(1, 1000, 1000 // num_steps) + num_timesteps = 1000 + ratio = (num_timesteps - 1) / (num_steps - 1) + timesteps = (np.arange(0, num_steps) * ratio).round().astype(np.int64) + alphas, alphas_prev = self._get_initial_alphas(timesteps) progbar = keras.utils.Progbar(len(timesteps)) iteration = 0