add custom handler

handler.py

@@ -2,6 +2,7 @@ from typing import Dict, List, Any
 import base64
 
 import tensorflow as tf
+from tensorflow import keras
 from keras_cv.models.generative.stable_diffusion.text_encoder import TextEncoder
 from keras_cv.models.generative.stable_diffusion.clip_tokenizer import SimpleTokenizer
 from keras_cv.models.generative.stable_diffusion.constants import _UNCONDITIONAL_TOKENS
@@ -12,6 +13,11 @@ class EndpointHandler():
 
         self.tokenizer = SimpleTokenizer()
         self.text_encoder = TextEncoder(self.MAX_PROMPT_LENGTH)
+        text_encoder_weights_fpath = keras.utils.get_file(
+            origin="https://huggingface.co/fchollet/stable-diffusion/resolve/main/kcv_encoder.h5",
+            file_hash="4789e63e07c0e54d6a34a29b45ce81ece27060c499a709d556c7755b42bb0dc4",
+        )
+        self.text_encoder.load_weights(text_encoder_weights_fpath)        
         self.pos_ids = tf.convert_to_tensor([list(range(self.MAX_PROMPT_LENGTH))], dtype=tf.int32)    
 
     def _get_unconditional_context(self):