add support for textual conversion pipelines.

app.py

@@ -19,6 +19,7 @@ This Space lets you convert KerasCV Stable Diffusion weights to a format compati
     * [Traditional text2image fine-tuning](https://github.com/huggingface/diffusers/blob/main/examples/text_to_image): UNet
     
     **In case none of the `text_encoder_weights` and `unet_weights` is provided, nothing will be done.**
+* For Textual Inversion, you MUST provide a valid `placeholder_token` i.e., the text concept used for conducting Textual Inversion.
 * When providing the weights' links, ensure they're directly downloadable. Internally, the Space uses [`tf.keras.utils.get_file()`](https://www.tensorflow.org/api_docs/python/tf/keras/utils/get_file) to retrieve the weights locally. 
 * If you don't provide `your_hf_token` the converted pipeline won't be pushed. 
 
@@ -26,7 +27,7 @@ Check [here](https://github.com/huggingface/diffusers/blob/31be42209ddfdb69d9640
 """
 
 
-def run(hf_token, text_encoder_weights, unet_weights, repo_prefix):
+def run(hf_token, text_encoder_weights, unet_weights, placeholder_token, repo_prefix):
     if text_encoder_weights == "":
         text_encoder_weights = None
     if unet_weights == "":
@@ -35,7 +36,12 @@ def run(hf_token, text_encoder_weights, unet_weights, repo_prefix):
     if text_encoder_weights is None and unet_weights is None:
         return "❌ No fine-tuned weights provided, nothing to do."
 
-    pipeline = run_conversion(text_encoder_weights, unet_weights)
+    if placeholder_token == "":
+        placeholder_token = None
+    if placeholder_token is not None and text_encoder_weights is None:
+        return "❌ Placeholder token provided but no text encoder weights were provided. Cannot proceed."
+
+    pipeline = run_conversion(text_encoder_weights, unet_weights, placeholder_token)
     output_path = "kerascv_sd_diffusers_pipeline"
     pipeline.save_pretrained(output_path)
 
@@ -48,6 +54,7 @@ def run(hf_token, text_encoder_weights, unet_weights, repo_prefix):
         base_model=PRETRAINED_CKPT,
         repo_folder=output_path,
         weight_paths=weight_paths,
+        placeholder_token=placeholder_token,
     )
     push_str = push_to_hub(hf_token, output_path, repo_prefix)
     return push_str
@@ -61,6 +68,7 @@ demo = gr.Interface(
         gr.Text(max_lines=1, label="your_hf_token"),
         gr.Text(max_lines=1, label="text_encoder_weights"),
         gr.Text(max_lines=1, label="unet_weights"),
+        gr.Text(max_lines=1, label="placeholder_token"),
         gr.Text(max_lines=1, label="output_repo_prefix"),
     ],
     outputs=[gr.Markdown(label="output")],

convert.py

@@ -4,7 +4,7 @@ from diffusers import (AutoencoderKL, StableDiffusionPipeline,
                        UNet2DConditionModel)
 from diffusers.pipelines.stable_diffusion.safety_checker import \
     StableDiffusionSafetyChecker
-from transformers import CLIPTextModel
+from transformers import CLIPTextModel, CLIPTokenizer
 
 from conversion_utils import populate_text_encoder, populate_unet
 
@@ -21,6 +21,7 @@ def initialize_pt_models():
     pt_text_encoder = CLIPTextModel.from_pretrained(
         PRETRAINED_CKPT, subfolder="text_encoder", revision=REVISION
     )
+    pt_tokenizer = CLIPTokenizer.from_pretrained(PRETRAINED_CKPT, subfolder="tokenizer")
     pt_vae = AutoencoderKL.from_pretrained(
         PRETRAINED_CKPT, subfolder="vae", revision=REVISION
     )
@@ -31,7 +32,7 @@ def initialize_pt_models():
         PRETRAINED_CKPT, subfolder="safety_checker", revision=NON_EMA_REVISION
     )
 
-    return pt_text_encoder, pt_vae, pt_unet, pt_safety_checker
+    return pt_text_encoder, pt_tokenizer, pt_vae, pt_unet, pt_safety_checker
 
 
 def initialize_tf_models(text_encoder_weights: str, unet_weights: str):
@@ -56,8 +57,18 @@ def initialize_tf_models(text_encoder_weights: str, unet_weights: str):
     return tf_sd_model, tf_text_encoder, tf_vae, tf_unet
 
 
-def run_conversion(text_encoder_weights: str = None, unet_weights: str = None):
-    pt_text_encoder, pt_vae, pt_unet, pt_safety_checker = initialize_pt_models()
+def run_conversion(
+    text_encoder_weights: str = None,
+    unet_weights: str = None,
+    placeholder_token: str = None,
+):
+    (
+        pt_text_encoder,
+        pt_tokenizer,
+        pt_vae,
+        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
     )
@@ -70,19 +81,30 @@ def run_conversion(text_encoder_weights: str = None, unet_weights: str = None):
         text_encoder_state_dict_from_tf = populate_text_encoder(tf_text_encoder)
         pt_text_encoder.load_state_dict(text_encoder_state_dict_from_tf)
         print("Populated PT text encoder from TF weights.")
+
     if unet_weights is not None:
         print("Loading fine-tuned UNet weights.")
         unet_weights_path = tf.keras.utils.get_file(origin=unet_weights)
         tf_unet.load_weights(unet_weights_path)
         unet_state_dict_from_tf = populate_unet(tf_unet)
         pt_unet.load_state_dict(unet_state_dict_from_tf)
-        print("Populated PT UNet from TF weights.")    
+        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,
         unet=pt_unet,
         text_encoder=pt_text_encoder,
+        tokenizer=pt_tokenizer,
         vae=pt_vae,
         safety_checker=pt_safety_checker,
         revision=None,

hub_utils/readme.py

@@ -3,7 +3,12 @@ from typing import List
 
 
 # Copied from https://github.com/huggingface/diffusers/blob/31be42209ddfdb69d9640a777b32e9b5c6259bf0/examples/text_to_image/train_text_to_image_lora.py#L55
-def save_model_card(base_model=str, repo_folder=None, weight_paths: List = None):
+def save_model_card(
+    base_model=str,
+    repo_folder=None,
+    weight_paths: List = None,
+    placeholder_token: str = None,
+):
     yaml = f"""
 ---
 license: creativeml-openrail-m
@@ -26,5 +31,8 @@ The pipeline contained in this repository was created using [this Space](https:/
     if len(weight_paths) > 0:
         model_card += f"Following weight paths (KerasCV) were used \n: {weight_paths}"
 
+    if placeholder_token is not None:
+        model_card += "\nFollowing `placeholder_token` was added to the tokenizer: {placeholder_token}."
+
     with open(os.path.join(repo_folder, "README.md"), "w") as f:
         f.write(yaml + model_card)