Fix LoRA offload detection for alignment hooks (#10914)

src/diffusers/loaders/lora_base.py

@@ -462,14 +462,16 @@ def _func_optionally_disable_offloading(_pipeline):
             if not isinstance(component, nn.Module):
                 continue
             is_group_offload = is_group_offload or _is_group_offload_enabled(component)
-            if not hasattr(component, "_hf_hook"):
-                continue
-            is_model_cpu_offload = is_model_cpu_offload or isinstance(component._hf_hook, CpuOffload)
-            is_sequential_cpu_offload = is_sequential_cpu_offload or (
-                isinstance(component._hf_hook, AlignDevicesHook)
-                or hasattr(component._hf_hook, "hooks")
-                and isinstance(component._hf_hook.hooks[0], AlignDevicesHook)
-            )
+            for module in component.modules():
+                if not hasattr(module, "_hf_hook"):
+                    continue
+                hooks = getattr(module._hf_hook, "hooks", (module._hf_hook,))
+                is_model_cpu_offload = is_model_cpu_offload or any(
+                    isinstance(hook, CpuOffload) for hook in hooks
+                )
+                is_sequential_cpu_offload = is_sequential_cpu_offload or any(
+                    isinstance(hook, AlignDevicesHook) and hook.offload for hook in hooks
+                )
 
         if is_sequential_cpu_offload or is_model_cpu_offload:
             logger.info(

tests/lora/test_lora_loader_utils.py

@@ -26,6 +26,7 @@
 from diffusers.loaders.peft import PeftAdapterMixin
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.utils.import_utils import is_peft_available
+from diffusers.utils.testing_utils import require_accelerate
 
 from ..testing_utils import CaptureLogger, require_peft_backend
 
@@ -118,6 +119,35 @@ def test_local_directory_with_multiple_files_warns_and_uses_first(tmp_path, monk
     assert "contains more than one weights file" in cap_logger.out
 
 
+@require_accelerate
+def test_alignment_hook_does_not_enable_sequential_cpu_offload():
+    from accelerate.hooks import AlignDevicesHook, add_hook_to_module
+
+    component = torch.nn.Linear(1, 1)
+    add_hook_to_module(component, AlignDevicesHook(execution_device="cpu"))
+    pipeline = Mock(hf_device_map=None, components={"model": component})
+
+    offload_state = lora_base._func_optionally_disable_offloading(pipeline)
+
+    assert offload_state == (False, False, False)
+    assert hasattr(component, "_hf_hook")
+
+
+@require_accelerate
+def test_sequential_cpu_offload_is_detected_and_disabled():
+    from accelerate import cpu_offload
+
+    component = torch.nn.Sequential(torch.nn.Linear(1, 1))
+    cpu_offload(component, execution_device=torch.device("cuda"))
+    pipeline = Mock(hf_device_map=None, components={"model": component})
+
+    offload_state = lora_base._func_optionally_disable_offloading(pipeline)
+
+    assert offload_state == (False, True, False)
+    assert not hasattr(component, "_hf_hook")
+    assert not hasattr(component[0], "_hf_hook")
+
+
 @require_peft_backend
 def test_unfuse_lora_partial_components_keeps_merged_adapters_in_sync():
     """Regression test for gh-14214.