improve runtime.

app.py

@@ -27,7 +27,7 @@ def run(hf_token, text_encoder_weights, unet_weights, repo_prefix):
         text_encoder_weights = None
     if unet_weights == "":
         unet_weights = None
-    print(f"unet_weights: {unet_weights}")
+
     pipeline = run_conversion(text_encoder_weights, unet_weights)
     output_path = "kerascv_sd_diffusers_pipeline"
     pipeline.save_pretrained(output_path)

convert.py

@@ -13,6 +13,7 @@ PRETRAINED_CKPT = "CompVis/stable-diffusion-v1-4"
 REVISION = None
 NON_EMA_REVISION = None
 IMG_HEIGHT = IMG_WIDTH = 512
+MAX_SEQ_LENGTH = 77
 
 
 def initialize_pt_models():
@@ -34,17 +35,25 @@ def initialize_pt_models():
     return pt_text_encoder, pt_vae, pt_unet, pt_safety_checker
 
 
-def initialize_tf_models():
+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."""
     tf_sd_model = keras_cv.models.StableDiffusion(
         img_height=IMG_HEIGHT, img_width=IMG_WIDTH
     )
-    _ = tf_sd_model.text_to_image("Cartoon")  # To download the weights.
-
-    tf_text_encoder = tf_sd_model.text_encoder
+    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
-    tf_unet = tf_sd_model.diffusion_model
+    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
 
 
@@ -55,11 +64,11 @@ def run_conversion(text_encoder_weights: str = None, unet_weights: str = None):
 
     if text_encoder_weights is not None:
         print("Loading fine-tuned text encoder weights.")
-        text_encoder_weights_path = tf.keras.utils.get_file(text_encoder_weights)
+        text_encoder_weights_path = tf.keras.utils.get_file(origin=text_encoder_weights)
         tf_text_encoder.load_weights(text_encoder_weights_path)
     if unet_weights is not None:
         print("Loading fine-tuned UNet weights.")
-        unet_weights_path = tf.keras.utils.get_file(unet_weights)
+        unet_weights_path = tf.keras.utils.get_file(origin=unet_weights)
         tf_unet.load_weights(unet_weights_path)
 
     text_encoder_state_dict_from_tf = populate_text_encoder(tf_text_encoder)