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

Fix offload issues with lora

Browse Source
This commit is contained in:
Disty0
2024-11-30 00:16:38 +03:00
parent 797ad1f20f
commit 881fa1183c
2 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/lora/networks.py
+11 -3
View File
@@ -280,8 +280,6 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No
if len(loaded_networks) > 0:
devices.torch_gc()
if shared.opts.diffusers_offload_mode == "balanced":
sd_models.apply_balanced_offload(shared.sd_model)
t1 = time.time()
timer['load'] = t1 - t0
@@ -375,7 +373,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
if module is not None and hasattr(self, 'weight'):
try:
with devices.inference_context():
weight = self.weight # calculate quant weights once
weight = self.weight.to(devices.device) # calculate quant weights once
updown, ex_bias = module.calc_updown(weight)
if batch_updown is not None and updown is not None:
batch_updown += updown
@@ -385,6 +383,11 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
batch_ex_bias += ex_bias
else:
batch_ex_bias = ex_bias
if shared.opts.diffusers_offload_mode != "none":
if batch_updown is not None:
batch_updown = batch_updown.to(devices.cpu)
if batch_ex_bias is not None:
batch_ex_bias = batch_ex_bias.to(devices.cpu)
except RuntimeError as e:
extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
if debug:
@@ -408,6 +411,9 @@ def network_load(): # called from processing
timer['calc'] = 0
timer['apply'] = 0
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) # wrapped model compatiblility
if shared.opts.diffusers_offload_mode != "none":
sd_models.disable_offload(sd_model)
sd_models.move_model(sd_model, device=devices.cpu)
with pbar:
for component_name in ['text_encoder','text_encoder_2', 'unet', 'transformer']:
component = getattr(sd_model, component_name, None)
@@ -428,6 +434,8 @@ def network_load(): # called from processing
pbar.remove_task(task)
if debug:
shared.log.debug(f'Load network: type=LoRA component={component_name} modules={len(modules)} applied={applied}')
if shared.opts.diffusers_offload_mode != "none":
sd_models.set_diffuser_offload(sd_model, op="model")
if debug:
shared.log.debug(f'Load network: type=LoRA total={total_time():.2f} timers={timer}')
modules/sd_models.py
+9 -5
View File
@@ -405,7 +405,7 @@ def apply_balanced_offload(sd_model):
if hasattr(pipe, "_internal_dict"):
keys = pipe._internal_dict.keys() # pylint: disable=protected-access
else:
keys = get_signature(shared.sd_model).keys()
keys = get_signature(pipe).keys()
for module_name in keys: # pylint: disable=protected-access
module = getattr(pipe, module_name, None)
if isinstance(module, torch.nn.Module):
@@ -1448,10 +1448,14 @@ def disable_offload(sd_model):
from accelerate.hooks import remove_hook_from_module
if not getattr(sd_model, 'has_accelerate', False):
return
if hasattr(sd_model, 'components'):
for _name, model in sd_model.components.items():
if isinstance(model, torch.nn.Module):
remove_hook_from_module(model, recurse=True)
if hasattr(sd_model, "_internal_dict"):
keys = sd_model._internal_dict.keys() # pylint: disable=protected-access
else:
keys = get_signature(sd_model).keys()
for module_name in keys: # pylint: disable=protected-access
module = getattr(sd_model, module_name, None)
if isinstance(module, torch.nn.Module):
module = remove_hook_from_module(module, recurse=True)
sd_model.has_accelerate = False