From 71a4706d5fb3113ce6f225d5bcb29522d8a8f950 Mon Sep 17 00:00:00 2001 From: Marc Sun Date: Wed, 5 Feb 2025 18:09:00 +0100 Subject: [PATCH] fix --- src/diffusers/pipelines/pipeline_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index c4593b3e698b..77c580bf310d 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1042,7 +1042,7 @@ def remove_all_hooks(self): """ for _, model in self.components.items(): if isinstance(model, torch.nn.Module) and hasattr(model, "_hf_hook"): - accelerate.hooks.remove_hook_from_module(model, recurse=True) + accelerate.hooks.remove_hook_from_module(_unwrap_model(model), recurse=True) self._all_hooks = [] def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"):