AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

make enable_sequential_cpu_offload more generic for third-party devices (#4191)

Browse Source 
* make enable_sequential_cpu_offload more generic for third-party devices

* make style
This commit is contained in:
statelesshz
2023-07-21 19:45:09 +08:00
committed by GitHub
parent 1e853e240e
commit e2bbaa4f54
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/diffusers/pipelines/pipeline_utils.py
View File

@@ -1127,7 +1127,9 @@ class DiffusionPipeline(ConfigMixin):
if self.device.type != "cpu":
self.to("cpu", silence_dtype_warnings=True)
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
device_mod = getattr(torch, self.device.type, None)
if hasattr(device_mod, "empty_cache") and device_mod.is_available():
device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
for name, model in self.components.items():
if not isinstance(model, torch.nn.Module):