From b4dd7869770681270faf577294c8ea387281f014 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Chollet?= Date: Sun, 25 Sep 2022 09:51:08 -0700 Subject: [PATCH] Do image clipping in numpy space (outputs are already numpy arrays after `predict()`, no need to go back to TF then to numpy again) (#838) --- .../models/generative/stable_diffusion/stable_diffusion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/keras_cv/models/generative/stable_diffusion/stable_diffusion.py b/keras_cv/models/generative/stable_diffusion/stable_diffusion.py index 53b8729f05..0571ab9632 100644 --- a/keras_cv/models/generative/stable_diffusion/stable_diffusion.py +++ b/keras_cv/models/generative/stable_diffusion/stable_diffusion.py @@ -23,6 +23,7 @@ import math +import numpy as np import tensorflow as tf from tensorflow import keras @@ -168,7 +169,7 @@ def text_to_image( # Decoding stage decoded = self.decoder.predict_on_batch(latent) decoded = ((decoded + 1) / 2) * 255 - return tf.cast(tf.clip_by_value(decoded, 0, 255), dtype=tf.uint8) + return np.clip(decoded, 0, 255).astype("uint8") def _get_timestep_embedding(self, timestep, batch_size, dim=320, max_period=10000): half = dim // 2