lora: only remove hooks that we add back

yiyixuxu · yiyixuxu · commit d12531ddf7ac · 2025-06-22T12:32:04.000+02:00
diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py
@@ -456,7 +456,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)
 

Original file line number	Diff line number	Diff line change
`@@ -456,7 +456,8 @@ def _func_optionally_disable_offloading(_pipeline):`
`456`	`456`	`logger.info(`
`457`	`457`	"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."
`458`	`458`	`)`
`459`		`- remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
	`459`	`+ if is_sequential_cpu_offload or is_model_cpu_offload:`
	`460`	`+ remove_hook_from_module(component, recurse=is_sequential_cpu_offload)`
`460`	`461`
`461`	`462`	`return (is_model_cpu_offload, is_sequential_cpu_offload)`
`462`	`463`