Compare commits

...

3 Commits

Author SHA1 Message Date
Marc Sun
72000d6b72 Merge remote-tracking branch 'upstream/main' into fix-hook-removal-with-warped-model 2025-03-04 10:50:32 +01:00
Dhruv Nair
d50a96c81f Merge branch 'main' into fix-hook-removal-with-warped-model 2025-02-12 11:11:40 +05:30
Marc Sun
71a4706d5f fix 2025-02-05 18:09:00 +01:00

View File

@@ -1069,7 +1069,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
"""
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"):