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

improve offloading

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-12-25 10:24:02 +00:00
parent 76aae17e6f
commit 0b1e6d2d3c
8 changed files with 47 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_offload.py
+36 -8
View File
@@ -18,7 +18,7 @@ offload_warn = ['sc', 'sd3', 'f1', 'f2', 'h1', 'hunyuandit', 'auraflow', 'omnige
offload_post = ['h1']
offload_hook_instance = None
balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline']
no_split_module_classes = ["Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "WanTransformerBlock"]
no_split_module_classes = ["Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "WanTransformerBlock", "ModuleDict", "ModuleList"]
accelerate_dtype_byte_size = None
move_stream = None
@@ -240,14 +240,26 @@ class OffloadHook(accelerate.hooks.ModelHook):
max_memory = { device_index: self.gpu, "cpu": self.cpu }
device_map = getattr(module, "balanced_offload_device_map", None)
if (device_map is None) or (max_memory != getattr(module, "balanced_offload_max_memory", None)):
device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory, no_split_module_classes=no_split_module_classes, verbose=verbose)
device_map = accelerate.infer_auto_device_map(module,
max_memory=max_memory,
no_split_module_classes=no_split_module_classes,
verbose=verbose,
clean_result=False,
)
offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
if devices.backend == "directml":
for k, v in device_map.items():
if isinstance(v, int):
device_map[k] = f"{devices.device.type}:{v}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
if debug:
shared.log.trace(f'Offload: type=balanced op=dispatch map={device_map}')
if device_map is not None:
module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
module = accelerate.dispatch_model(module,
main_device=torch.device(devices.device),
device_map=device_map,
offload_dir=offload_dir,
force_hooks=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
@@ -291,6 +303,15 @@ def get_pipe_variants(pipe=None):
def get_module_names(pipe=None, exclude=None):
def is_valid(module):
if isinstance(getattr(pipe, module, None), torch.nn.ModuleDict):
return True
if isinstance(getattr(pipe, module, None), torch.nn.ModuleList):
return True
if isinstance(getattr(pipe, module, None), torch.nn.Module):
return True
return False
if exclude is None:
exclude = []
if pipe is None:
@@ -298,12 +319,19 @@ def get_module_names(pipe=None, exclude=None):
pipe = shared.sd_model
else:
return []
if hasattr(pipe, "_internal_dict"):
modules_names = pipe._internal_dict.keys() # pylint: disable=protected-access
else:
modules_names = get_signature(pipe).keys()
modules_names = []
try:
dict_keys = pipe._internal_dict.keys() # pylint: disable=protected-access
modules_names.extend(dict_keys)
except Exception:
pass
try:
dict_keys = get_signature(pipe).keys()
modules_names.extend(dict_keys)
except Exception:
pass
modules_names = [m for m in modules_names if m not in exclude and not m.startswith('_')]
modules_names = [m for m in modules_names if isinstance(getattr(pipe, m, None), torch.nn.Module)]
modules_names = [m for m in modules_names if is_valid(m)]
modules_names = sorted(set(modules_names))
return modules_names