Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

set offload_buffers to true on balanced offload

Browse Source
This commit is contained in:
Disty0
2025-09-28 16:23:29 +03:00
parent 3275535867
commit 02f03d56a7
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_offload.py
+1 -1
View File
@@ -245,7 +245,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
if isinstance(device_map[v], int):
device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
if device_map is not None:
module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir, offload_buffers=True)
module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
module.balanced_offload_device_map = device_map
module.balanced_offload_max_memory = max_memory