better conditioning on weight porting.

convert.py

@@ -6,8 +6,7 @@ from diffusers.pipelines.stable_diffusion.safety_checker import \
     StableDiffusionSafetyChecker
 from transformers import CLIPTextModel
 
-from conversion_utils import (populate_text_encoder, populate_unet,
-                              run_assertion)
+from conversion_utils import populate_text_encoder, populate_unet
 
 PRETRAINED_CKPT = "CompVis/stable-diffusion-v1-4"
 REVISION = None
@@ -68,31 +67,18 @@ def run_conversion(text_encoder_weights: str = None, unet_weights: str = None):
         print("Loading fine-tuned 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)
+        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.")    
 
-    text_encoder_state_dict_from_tf = populate_text_encoder(tf_text_encoder)
-    unet_state_dict_from_tf = populate_unet(tf_unet)
-    print("Conversion done, now running optional assertions...")
-
-    # Since we cannot compare the fine-tuned weights.
-    if text_encoder_weights is None:
-        text_encoder_state_dict_from_pt = pt_text_encoder.state_dict()
-        run_assertion(text_encoder_state_dict_from_pt, text_encoder_state_dict_from_tf)
-    if unet_weights is None:
-        unet_state_dict_from_pt = pt_unet.state_dict()
-        run_assertion(unet_state_dict_from_pt, unet_state_dict_from_tf)
-
-    if text_encoder_weights is None or unet_weights is None:
-        print(
-            "Assertions successful, populating the converted parameters into the diffusers models..."
-        )
-    pt_text_encoder.load_state_dict(text_encoder_state_dict_from_tf)
-    pt_unet.load_state_dict(unet_state_dict_from_tf)
-
-    print("Parameters ported, preparing StabelDiffusionPipeline...")
+    print("Weights ported, preparing StabelDiffusionPipeline...")
     pipeline = StableDiffusionPipeline.from_pretrained(
         PRETRAINED_CKPT,
         unet=pt_unet,