|
|
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} |
|
|
|
