Fix the issue (keras-team#2102): Add cast to float32 from float16 int…

…o Stable Diffusion
y-vectorfield · Nov 2, 2023 · 3c2f897 · 3c2f897
1 parent 832e2d9
commit 3c2f897
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/keras_cv/models/stable_diffusion/stable_diffusion.py b/keras_cv/models/stable_diffusion/stable_diffusion.py
@@ -235,6 +235,7 @@ def generate_image(
                 + unconditional_guidance_scale * (latent - unconditional_latent)
             )
             a_t, a_prev = alphas[index], alphas_prev[index]
+            latent = ops.cast(latent, latent_prev.dtype)
             pred_x0 = (latent_prev - math.sqrt(1 - a_t) * latent) / math.sqrt(
                 a_t
             )