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

jumbo update: add wan22-vace, improve offloading, add offloading-using-streams, change logging-levels, refactor some default packages

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-09-17 12:32:27 -04:00
parent 6a021e7743
commit cc6101ecb2
17 changed files with 156 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_offload.py
+29 -9
View File
@@ -18,6 +18,7 @@ offload_post = ['h1']
offload_hook_instance = None
balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline']
accelerate_dtype_byte_size = None
move_stream = None
def dtype_byte_size(dtype: torch.dtype):
@@ -176,7 +177,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
self.last_post = None
self.last_cls = None
gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}'
shared.log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never} pre={shared.opts.diffusers_offload_pre}')
shared.log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never} pre={shared.opts.diffusers_offload_pre} streams={shared.opts.diffusers_offload_streams}')
self.validate()
super().__init__()
@@ -210,8 +211,12 @@ class OffloadHook(accelerate.hooks.ModelHook):
def pre_forward(self, module, *args, **kwargs):
_id = id(module)
if (self.last_pre != _id) and (module.__class__.__name__ != self.last_cls) and self.offload_allowed(module): # offload every other module first time when new module starts pre-forward
do_offload = (self.last_pre != _id) or (module.__class__.__name__ != self.last_cls)
if do_offload and self.offload_allowed(module): # offload every other module first time when new module starts pre-forward
if shared.opts.diffusers_offload_pre:
t0 = time.time()
debug_move(f'Offload: type=balanced op=pre module={module.__class__.__name__}')
for pipe in get_pipe_variants():
for module_name in get_module_names(pipe):
@@ -220,15 +225,16 @@ class OffloadHook(accelerate.hooks.ModelHook):
if (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
apply_balanced_offload_to_module(module_instance, op='pre')
self.last_cls = module.__class__.__name__
self.last_pre = _id
process_timer.add('offload', time.time() - t0)
if not devices.same_device(module.device, devices.device): # move-to-device
t0 = time.time()
device_index = torch.device(devices.device).index
if device_index is None:
device_index = 0
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):
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)
offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
if devices.backend == "directml":
@@ -241,13 +247,15 @@ class OffloadHook(accelerate.hooks.ModelHook):
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
process_timer.add('onload', time.time() - t0)
if debug:
for pipe in get_pipe_variants():
for _i, pipe in enumerate(get_pipe_variants()):
for module_name in get_module_names(pipe):
module_instance = getattr(pipe, module_name, None)
shared.log.trace(f'Offload: type=balanced op=pre:status module={module_instance.__class__.__name__} device={module_instance.device} dtype={module_instance.dtype}')
shared.log.trace(f'Offload: type=balanced op=pre:status forward={module.__class__.__name__} module={module_name} class={module_instance.__class__.__name__} pipe={_i} device={module_instance.device} dtype={module_instance.dtype}')
self.last_pre = _id
return args, kwargs
def post_forward(self, module, output):
@@ -283,6 +291,7 @@ def get_module_names(pipe=None, exclude=[]):
modules_names = get_signature(pipe).keys()
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 = list(sorted(set(modules_names)))
return modules_names
@@ -308,6 +317,17 @@ def get_module_sizes(pipe=None, exclude=[]):
def move_module_to_cpu(module, op='unk', force:bool=False):
def do_move(module):
if shared.opts.diffusers_offload_streams:
global move_stream # pylint: disable=global-statement
if move_stream is None:
move_stream = torch.cuda.Stream(device=devices.device)
with torch.cuda.stream(move_stream):
module = module.to(devices.cpu)
else:
module = module.to(devices.cpu)
return module
try:
module_name = getattr(module, "module_name", module.__class__.__name__)
module_size = offload_hook_instance.offload_map.get(module_name, offload_hook_instance.model_size())
@@ -318,17 +338,17 @@ def move_module_to_cpu(module, op='unk', force:bool=False):
op = f'{op}:skip'
if force:
op = f'{op}:force'
module = module.to(devices.cpu)
module = do_move(module)
used_gpu -= module_size
elif module_cls in offload_hook_instance.offload_never:
op = f'{op}:never'
elif module_cls in offload_hook_instance.offload_always:
op = f'{op}:always'
module = module.to(devices.cpu)
module = do_move(module)
used_gpu -= module_size
elif perc_gpu > shared.opts.diffusers_offload_min_gpu_memory:
op = f'{op}:mem'
module = module.to(devices.cpu)
module = do_move(module)
used_gpu -= module_size
if debug:
quant = getattr(module, "quantization_method", None)