fix: textual inversion utility.

convert.py

@@ -35,26 +35,51 @@ def initialize_pt_models():
     return pt_text_encoder, pt_tokenizer, pt_vae, pt_unet, pt_safety_checker
 
 
-def initialize_tf_models(text_encoder_weights: str, unet_weights: str):
-    """Initializes the separate models of Stable Diffusion from KerasCV and downloads
-    their pre-trained weights."""
+def initialize_tf_models(
+    text_encoder_weights: str, unet_weights: str, placeholder_token: str = None
+):
+    """Initializes the separate models of Stable Diffusion from KerasCV and optionally
+    downloads their pre-trained weights."""
     tf_sd_model = keras_cv.models.StableDiffusion(
         img_height=IMG_HEIGHT, img_width=IMG_WIDTH
     )
+
     if text_encoder_weights is None:
         tf_text_encoder = tf_sd_model.text_encoder
     else:
         tf_text_encoder = keras_cv.models.stable_diffusion.TextEncoder(
             MAX_SEQ_LENGTH, download_weights=False
         )
-    tf_vae = tf_sd_model.image_encoder
+
     if unet_weights is None:
         tf_unet = tf_sd_model.diffusion_model
     else:
         tf_unet = keras_cv.models.stable_diffusion.DiffusionModel(
             IMG_HEIGHT, IMG_WIDTH, MAX_SEQ_LENGTH, download_weights=False
         )
-    return tf_sd_model, tf_text_encoder, tf_vae, tf_unet
+
+    tf_tokenizer = tf_sd_model.tokenizer
+    if placeholder_token is not None:
+        tf_tokenizer.add_tokens(placeholder_token)
+
+    return tf_text_encoder, tf_unet, tf_tokenizer
+
+
+def create_new_text_encoder(tf_text_encoder, tf_tokenizer):
+    """Initializes a fresh text encoder in case the weights are from Textual Inversion.
+
+    Reference: https://keras.io/examples/generative/fine_tune_via_textual_inversion/
+    """
+    new_vocab_size = len(tf_tokenizer.vocab)
+    new_text_encoder = keras_cv.models.stable_diffusion.TextEncoder(
+        MAX_SEQ_LENGTH,
+        vocab_size=new_vocab_size,
+        download_weights=False,
+    )
+
+    old_position_weights = tf_text_encoder.layers[2].position_embedding.get_weights()
+    new_text_encoder.layers[2].position_embedding.set_weights(old_position_weights)
+    return new_text_encoder
 
 
 def run_conversion(
@@ -69,11 +94,23 @@ def run_conversion(
         pt_unet,
         pt_safety_checker,
     ) = initialize_pt_models()
-    tf_sd_model, tf_text_encoder, tf_vae, tf_unet = initialize_tf_models(
-        text_encoder_weights, unet_weights
+    tf_text_encoder, tf_unet, tf_tokenizer = initialize_tf_models(
+        text_encoder_weights, unet_weights, placeholder_token
     )
     print("Pre-trained model weights downloaded.")
 
+    if placeholder_token is not None:
+        print("Initializing a new text encoder with the placeholder token...")
+        tf_text_encoder = create_new_text_encoder(tf_text_encoder, tf_tokenizer)
+
+        print("Adding the placeholder token to PT CLIPTokenizer...")
+        num_added_tokens = pt_tokenizer.add_tokens(placeholder_token)
+        if num_added_tokens == 0:
+            raise ValueError(
+                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
+                " `placeholder_token` that is not already in the tokenizer."
+            )
+
     if text_encoder_weights is not None:
         print("Loading fine-tuned text encoder weights.")
         text_encoder_weights_path = tf.keras.utils.get_file(origin=text_encoder_weights)
@@ -90,15 +127,6 @@ def run_conversion(
         pt_unet.load_state_dict(unet_state_dict_from_tf)
         print("Populated PT UNet from TF weights.")
 
-    if placeholder_token is not None:
-        print("Adding the placeholder_token to CLIPTokenizer...")
-        num_added_tokens = pt_tokenizer.add_tokens(placeholder_token)
-        if num_added_tokens == 0:
-            raise ValueError(
-                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
-                " `placeholder_token` that is not already in the tokenizer."
-            )
-
     print("Weights ported, preparing StabelDiffusionPipeline...")
     pipeline = StableDiffusionPipeline.from_pretrained(
         PRETRAINED_CKPT,