Update handler.py

handler.py

@@ -8,17 +8,21 @@ from keras_cv.models.generative.stable_diffusion.decoder import Decoder
 
 class EndpointHandler():
     def __init__(self, path=""):
-      img_height = 512
-      img_width = 512
-      img_height = round(img_height / 128) * 128
-      img_width = round(img_width / 128) * 128
-
-      self.decoder = Decoder(img_height, img_width)
-      decoder_weights_fpath = keras.utils.get_file(
-          origin="https://huggingface.co/fchollet/stable-diffusion/resolve/main/kcv_decoder.h5",
-          file_hash="ad350a65cc8bc4a80c8103367e039a3329b4231c2469a1093869a345f55b1962",
-      )
-      self.decoder.load_weights(decoder_weights_fpath)
+        gpus = tf.config.list_physical_devices(device_type = 'GPU')
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        
+        img_height = 512
+        img_width = 512
+        img_height = round(img_height / 128) * 128
+        img_width = round(img_width / 128) * 128
+        
+        self.decoder = Decoder(img_height, img_width)
+        decoder_weights_fpath = keras.utils.get_file(
+            origin="https://huggingface.co/fchollet/stable-diffusion/resolve/main/kcv_decoder.h5",
+            file_hash="ad350a65cc8bc4a80c8103367e039a3329b4231c2469a1093869a345f55b1962",
+        )
+        self.decoder.load_weights(decoder_weights_fpath)
 
     def __call__(self, data: Dict[str, Any]) -> str:
         # get inputs