From 02f03d56a71b2a579af51293dcd7f550fe7dab0a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 28 Sep 2025 16:23:29 +0300 Subject: [PATCH] set offload_buffers to true on balanced offload --- modules/sd_offload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 19e59de3d..4126b81e8 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -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