[lora] only remove hooks that we add back (#11768)

yiyixuxu · web-flow · commit 7bc0a07b1947 · 2025-06-23T16:49:19.000-10:00
up
diff --git 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
@@ -427,7 +427,8 @@ def load_textual_inversion(
                         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

Original file line number	Diff line number	Diff line change
`@@ -453,7 +453,8 @@ def _func_optionally_disable_offloading(_pipeline):`
`453`	`453`	`logger.info(`
`454`	`454`	"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."
`455`	`455`	`)`
`456`		`- remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
	`456`	`+ if is_sequential_cpu_offload or is_model_cpu_offload:`
	`457`	`+ remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
`457`	`458`
`458`	`459`	`return (is_model_cpu_offload, is_sequential_cpu_offload)`
`459`	`460`
Original file line number	Diff line number	Diff line change
`@@ -427,7 +427,8 @@ def load_textual_inversion(`
`427`	`427`	`logger.info(`
`428`	`428`	"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."
`429`	`429`	`)`
`430`		`- remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
	`430`	`+ if is_sequential_cpu_offload or is_model_cpu_offload:`
	`431`	`+ remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
`431`	`432`
`432`	`433`	`# 7.2 save expected device and dtype`
`433`	`434`	`device = text_encoder.device`