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

…o Stable Diffusion
y-vectorfield · Oct 26, 2023 · cc214e6 · cc214e6
1 parent 832e2d9
commit cc214e6
Showing 1 changed file with 2 additions 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,8 @@ def generate_image(
                 + unconditional_guidance_scale * (latent - unconditional_latent)
             )
             a_t, a_prev = alphas[index], alphas_prev[index]
+            if latent.dtype == "float16":
+                latent = ops.cast(latent, "float32")
             pred_x0 = (latent_prev - math.sqrt(1 - a_t) * latent) / math.sqrt(
                 a_t
             )