From dcce89044898e5f400723cb09fd59dd3f9984ee1 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 20 Dec 2024 01:49:16 -0600 Subject: [PATCH] Fix `save_lora_weights` in `pipeline_utils.py` (#1643) --- optimum/habana/diffusers/pipelines/pipeline_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/optimum/habana/diffusers/pipelines/pipeline_utils.py b/optimum/habana/diffusers/pipelines/pipeline_utils.py index 2cbfd57d2b..efd1b36693 100644 --- a/optimum/habana/diffusers/pipelines/pipeline_utils.py +++ b/optimum/habana/diffusers/pipelines/pipeline_utils.py @@ -28,10 +28,10 @@ from diffusers.utils.torch_utils import is_compiled_module from huggingface_hub import create_repo -from optimum.habana.utils import to_device_dtype from optimum.utils import logging from ...transformers.gaudi_configuration import GaudiConfig +from ...utils import to_device_dtype logger = logging.get_logger(__name__) @@ -396,7 +396,8 @@ def save_lora_weights( text_encoder_2_lora_layers = to_device_dtype(text_encoder_2_lora_layers, target_device=torch.device("cpu")) # text_encoder_2_lora_layers is only supported by some diffuser pipelines - if text_encoder_2_lora_layers: + signature = inspect.signature(super().save_lora_weights) + if "text_encoder_2_lora_layers" in signature.parameters: return super().save_lora_weights( save_directory, unet_lora_layers,