diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 26bd938b2734..eeff1e9e0b74 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1069,7 +1069,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"):