update custom handler

handler.py

@@ -1,27 +1,41 @@
 from typing import Dict, List, Any
 import base64
 
-import logging
-
 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
+from keras_cv.models.stable_diffusion.text_encoder import TextEncoder
+from keras_cv.models.stable_diffusion.text_encoder import TextEncoderV2
+from keras_cv.models.stable_diffusion.clip_tokenizer import SimpleTokenizer
+from keras_cv.models.stable_diffusion.constants import _UNCONDITIONAL_TOKENS
 
 class EndpointHandler():
-    def __init__(self, path=""):
+    def __init__(self, path="", version="2"):
         self.MAX_PROMPT_LENGTH = 77
 
         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.text_encoder = self._instantiate_text_encoder(version)
         self.pos_ids = tf.convert_to_tensor([list(range(self.MAX_PROMPT_LENGTH))], dtype=tf.int32)    
 
+    def _instantiate_text_encoder(self, version: str):
+        if version == "1.4":
+            text_encoder_weights_fpath = keras.utils.get_file(
+                origin="https://huggingface.co/fchollet/stable-diffusion/resolve/main/kcv_encoder.h5",
+                file_hash="4789e63e07c0e54d6a34a29b45ce81ece27060c499a709d556c7755b42bb0dc4",
+            )
+            text_encoder = TextEncoder(self.MAX_PROMPT_LENGTH)
+            text_encoder.load_weights(text_encoder_weights_fpath)
+            return text_encoder
+        elif version == "2":
+            text_encoder_weights_fpath = keras.utils.get_file(
+                origin="https://huggingface.co/ianstenbit/keras-sd2.1/resolve/main/text_encoder_v2_1.h5",
+                file_hash="985002e68704e1c5c3549de332218e99c5b9b745db7171d5f31fcd9a6089f25b",
+            )
+            text_encoder = TextEncoderV2(self.MAX_PROMPT_LENGTH)
+            text_encoder.load_weights(text_encoder_weights_fpath)
+            return text_encoder
+        else:
+            return f"v{version} is not supported"
+
     def _get_unconditional_context(self):
         unconditional_tokens = tf.convert_to_tensor(
             [_UNCONDITIONAL_TOKENS], dtype=tf.int32

requirements.txt

@@ -1,3 +1,3 @@
-keras-cv
-tensorflow
+keras-cv==0.4
+tensorflow==2.11
 tensorflow_datasets