| from typing import Dict, List, Any |
| import sys |
| import base64 |
|
|
| import tensorflow as tf |
| from tensorflow import keras |
| 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="", version="2"): |
| self.MAX_PROMPT_LENGTH = 77 |
|
|
| self.text_encoder = self._instantiate_text_encoder(version) |
| if isinstance(self.text_encoder, str): |
| sys.exit(self.text_encoder) |
|
|
| self.tokenizer = SimpleTokenizer() |
| 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 |
| ) |
| unconditional_context = self.text_encoder.predict_on_batch( |
| [unconditional_tokens, self.pos_ids] |
| ) |
|
|
| return unconditional_context |
|
|
| def encode_text(self, prompt): |
| |
| inputs = self.tokenizer.encode(prompt) |
| if len(inputs) > self.MAX_PROMPT_LENGTH: |
| raise ValueError( |
| f"Prompt is too long (should be <= {self.MAX_PROMPT_LENGTH} tokens)" |
| ) |
| phrase = inputs + [49407] * (self.MAX_PROMPT_LENGTH - len(inputs)) |
| phrase = tf.convert_to_tensor([phrase], dtype=tf.int32) |
|
|
| context = self.text_encoder.predict_on_batch([phrase, self.pos_ids]) |
|
|
| return context |
|
|
| def get_contexts(self, encoded_text, batch_size): |
| encoded_text = tf.squeeze(encoded_text) |
| if encoded_text.shape.rank == 2: |
| encoded_text = tf.repeat( |
| tf.expand_dims(encoded_text, axis=0), batch_size, axis=0 |
| ) |
|
|
| context = encoded_text |
|
|
| unconditional_context = tf.repeat( |
| self._get_unconditional_context(), batch_size, axis=0 |
| ) |
|
|
| return context, unconditional_context |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| |
| prompt = data.pop("inputs", data) |
| batch_size = data.pop("batch_size", 1) |
|
|
| encoded_text = self.encode_text(prompt) |
| context, unconditional_context = self.get_contexts(encoded_text, batch_size) |
|
|
| context_b64 = base64.b64encode(context.numpy().tobytes()) |
| context_b64str = context_b64.decode() |
|
|
| unconditional_context_b64 = base64.b64encode(unconditional_context.numpy().tobytes()) |
| unconditional_context_b64str = unconditional_context_b64.decode() |
| |
| return {"context_b64str": context_b64str, "unconditional_context_b64str": unconditional_context_b64str} |
|
|
