From 7bc0a07b1947bcbe0e84bbe9ecf8ae2d234382c4 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 23 Jun 2025 16:49:19 -1000 Subject: [PATCH] [lora] only remove hooks that we add back (#11768) up --- src/diffusers/loaders/lora_base.py | 3 ++- src/diffusers/loaders/textual_inversion.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index 87b5b1a3fa..16f0d48365 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -453,7 +453,8 @@ def _func_optionally_disable_offloading(_pipeline): logger.info( "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again." ) - remove_hook_from_module(component, recurse=is_sequential_cpu_offload) + if is_sequential_cpu_offload or is_model_cpu_offload: + remove_hook_from_module(component, recurse=is_sequential_cpu_offload) return (is_model_cpu_offload, is_sequential_cpu_offload) diff --git a/src/diffusers/loaders/textual_inversion.py b/src/diffusers/loaders/textual_inversion.py index fb860b04a4..63fc97ed43 100644 --- a/src/diffusers/loaders/textual_inversion.py +++ b/src/diffusers/loaders/textual_inversion.py @@ -427,7 +427,8 @@ class TextualInversionLoaderMixin: logger.info( "Accelerate hooks detected. Since you have called `load_textual_inversion()`, the previous hooks will be first removed. Then the textual inversion parameters will be loaded and the hooks will be applied again." ) - remove_hook_from_module(component, recurse=is_sequential_cpu_offload) + if is_sequential_cpu_offload or is_model_cpu_offload: + remove_hook_from_module(component, recurse=is_sequential_cpu_offload) # 7.2 save expected device and dtype device = text_encoder.device