Update save lora weights for diffusers with text_encoder_2 layers (hu…

…ggingface#1626)
skavulya · Dec 19, 2024 · a51475f · a51475f
1 parent eee2d91
commit a51475f
Showing 1 changed file with 23 additions and 10 deletions.
diff --git a/optimum/habana/diffusers/pipelines/pipeline_utils.py b/optimum/habana/diffusers/pipelines/pipeline_utils.py
@@ -394,13 +394,26 @@ def save_lora_weights(
             text_encoder_lora_layers = to_device_dtype(text_encoder_lora_layers, target_device=torch.device("cpu"))
         if text_encoder_2_lora_layers:
             text_encoder_2_lora_layers = to_device_dtype(text_encoder_2_lora_layers, target_device=torch.device("cpu"))
-        return super().save_lora_weights(
-            save_directory,
-            unet_lora_layers,
-            text_encoder_lora_layers,
-            text_encoder_2_lora_layers,
-            is_main_process,
-            weight_name,
-            save_function,
-            safe_serialization,
-        )
+
+        # text_encoder_2_lora_layers is only supported by some diffuser pipelines
+        if text_encoder_2_lora_layers:
+            return super().save_lora_weights(
+                save_directory,
+                unet_lora_layers,
+                text_encoder_lora_layers,
+                text_encoder_2_lora_layers,
+                is_main_process,
+                weight_name,
+                save_function,
+                safe_serialization,
+            )
+        else:
+            return super().save_lora_weights(
+                save_directory,
+                unet_lora_layers,
+                text_encoder_lora_layers,
+                is_main_process,
+                weight_name,
+                save_function,
+                safe_serialization,
+            )