From fdb2ea8260c09c6d28e41e2a25bcfab6187ebfe9 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 22 Sep 2024 13:13:57 -0400 Subject: [PATCH] always keep lora on gpu --- CHANGELOG.md | 3 ++- .../Lora/extra_networks_lora.py | 2 +- extensions-builtin/Lora/network_full.py | 1 - extensions-builtin/Lora/network_ia3.py | 1 - extensions-builtin/Lora/network_lora.py | 2 +- extensions-builtin/Lora/network_oft.py | 4 ---- extensions-builtin/Lora/networks.py | 12 ----------- modules/processing_args.py | 1 + modules/processing_vae.py | 20 +++++++++++++++---- 9 files changed, 21 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 32da9f705..cd70c6d08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,7 +36,8 @@ - if lora contains no tags, lora name itself will be used as a tag - if prompt contains `_tags_` it will be used as placeholder for replacement, otherwise tags will be appended - used tags are also logged and registered in image metadata - - correct using of `extra_networks_default_multiplier` if not scale is specified + - correct using of `extra_networks_default_multiplier` if not scale is specified + - always keep lora on gpu - **text encoder**: - allow loading different custom text encoders: *clip-vit-l, clip-vit-g, t5* will automatically find appropriate encoder in the loaded model and replace it with loaded text encoder diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py index 2445d08e8..39ae11342 100644 --- a/extensions-builtin/Lora/extra_networks_lora.py +++ b/extensions-builtin/Lora/extra_networks_lora.py @@ -121,7 +121,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): self.active = True networks.originals.apply() # apply patches if networks.debug: - shared.log.debug("LoRA activate") + shared.log.debug(f"LoRA activate: networks={len(params_list)}") t1 = time.time() names, te_multipliers, unet_multipliers, dyn_dims = self.parse(p, params_list, step) networks.load_networks(names, te_multipliers, unet_multipliers, dyn_dims) diff --git a/extensions-builtin/Lora/network_full.py b/extensions-builtin/Lora/network_full.py index 346893706..273b26605 100644 --- a/extensions-builtin/Lora/network_full.py +++ b/extensions-builtin/Lora/network_full.py @@ -5,7 +5,6 @@ class ModuleTypeFull(network.ModuleType): def create_module(self, net: network.Network, weights: network.NetworkWeights): if all(x in weights.w for x in ["diff"]): return NetworkModuleFull(net, weights) - return None diff --git a/extensions-builtin/Lora/network_ia3.py b/extensions-builtin/Lora/network_ia3.py index 61c05ca02..6bf26b133 100644 --- a/extensions-builtin/Lora/network_ia3.py +++ b/extensions-builtin/Lora/network_ia3.py @@ -5,7 +5,6 @@ class ModuleTypeIa3(network.ModuleType): def create_module(self, net: network.Network, weights: network.NetworkWeights): if all(x in weights.w for x in ["weight"]): return NetworkModuleIa3(net, weights) - return None diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py index 6009002fd..dc54abe3f 100644 --- a/extensions-builtin/Lora/network_lora.py +++ b/extensions-builtin/Lora/network_lora.py @@ -47,7 +47,7 @@ class NetworkModuleLora(network.NetworkModule): if weight.shape != module.weight.shape: weight = weight.reshape(module.weight.shape) module.weight.copy_(weight) - module.to(device=devices.cpu, dtype=devices.dtype) + module = module.to(device=devices.device, dtype=devices.dtype) module.weight.requires_grad_(False) return module diff --git a/extensions-builtin/Lora/network_oft.py b/extensions-builtin/Lora/network_oft.py index 802492ca2..2005a62ec 100644 --- a/extensions-builtin/Lora/network_oft.py +++ b/extensions-builtin/Lora/network_oft.py @@ -8,19 +8,15 @@ class ModuleTypeOFT(network.ModuleType): def create_module(self, net: network.Network, weights: network.NetworkWeights): if all(x in weights.w for x in ["oft_blocks"]) or all(x in weights.w for x in ["oft_diag"]): return NetworkModuleOFT(net, weights) - return None # Supports both kohya-ss' implementation of COFT https://github.com/kohya-ss/sd-scripts/blob/main/networks/oft.py # and KohakuBlueleaf's implementation of OFT/COFT https://github.com/KohakuBlueleaf/LyCORIS/blob/dev/lycoris/modules/diag_oft.py class NetworkModuleOFT(network.NetworkModule): # pylint: disable=abstract-method def __init__(self, net: network.Network, weights: network.NetworkWeights): - super().__init__(net, weights) - self.lin_module = None self.org_module: list[torch.Module] = [self.sd_module] - self.scale = 1.0 # kohya-ss diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index a3a3ada92..16328af0a 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -421,15 +421,11 @@ def network_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]): def network_Linear_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.Linear_forward) network_apply_weights(self) return originals.Linear_forward(self, input) def network_QLinear_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.Linear_forward) network_apply_weights(self) return torch.nn.functional.linear(input, self.qweight, bias=self.bias) @@ -440,15 +436,11 @@ def network_Linear_load_state_dict(self, *args, **kwargs): def network_Conv2d_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.Conv2d_forward) network_apply_weights(self) return originals.Conv2d_forward(self, input) def network_QConv2d_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.Conv2d_forward) network_apply_weights(self) return self._conv_forward(input, self.qweight, self.bias) # pylint: disable=protected-access @@ -459,8 +451,6 @@ def network_Conv2d_load_state_dict(self, *args, **kwargs): def network_GroupNorm_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.GroupNorm_forward) network_apply_weights(self) return originals.GroupNorm_forward(self, input) @@ -471,8 +461,6 @@ def network_GroupNorm_load_state_dict(self, *args, **kwargs): def network_LayerNorm_forward(self, input): # pylint: disable=W0622 - # if shared.opts.lora_functional: - # return network_forward(self, input, originals.LayerNorm_forward) network_apply_weights(self) return originals.LayerNorm_forward(self, input) diff --git a/modules/processing_args.py b/modules/processing_args.py index ef7e9f30f..adb0992f2 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -268,6 +268,7 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 if isinstance(v, list) and len(v) > 0 and (isinstance(v[0], torch.Tensor) or isinstance(v[0], np.ndarray)): clean[k] = [x.shape for x in v] shared.log.debug(f'Diffuser pipeline: {model.__class__.__name__} task={sd_models.get_diffusers_task(model)} batch={p.iteration + 1}/{p.n_iter}x{p.batch_size} set={clean}') + if p.hdr_clamp or p.hdr_maximize or p.hdr_brightness != 0 or p.hdr_color != 0 or p.hdr_sharpen != 0: txt = 'HDR:' txt += f' Brightness={p.hdr_brightness}' if p.hdr_brightness != 0 else ' Brightness off' diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 45862da15..30f537632 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -3,7 +3,7 @@ import time import numpy as np import torch import torchvision.transforms.functional as TF -from modules import shared, devices, sd_models, sd_vae, sd_vae_taesd +from modules import shared, devices, sd_models, sd_vae, sd_vae_taesd, errors debug = os.environ.get('SD_VAE_DEBUG', None) is not None @@ -63,6 +63,8 @@ def full_vae_decode(latents, model): latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype) elif shared.opts.no_half_vae: latents = latents.to(torch.float32) + else: + latents = latents.to(model.vae.device) # normalize latents latents_mean = model.vae.config.get("latents_mean", None) @@ -77,8 +79,18 @@ def full_vae_decode(latents, model): latents = latents / scaling_factor if shift_factor: latents = latents + shift_factor - latents = latents.to(model.vae.device) - decoded = model.vae.decode(latents, return_dict=False)[0] + + vae_name = os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0] if sd_vae.loaded_vae_file is not None else "default" + vae_stats = f'name="{vae_name}" dtype={model.vae.dtype} device={model.vae.device} upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)}' + latents_stats = f'shape={latents.shape} dtype={latents.dtype} device={latents.device}' + stats = f'vae {vae_stats} latents {latents_stats}' + + try: + decoded = model.vae.decode(latents, return_dict=False)[0] + except Exception as e: + shared.log.error(f'VAE decode: {stats} {e}') + errors.display(e, 'VAE decode') + decoded = [] # delete vae after OpenVINO compile if 'VAE' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: @@ -94,7 +106,7 @@ def full_vae_decode(latents, model): t1 = time.time() if debug: log_debug(f'VAE memory: {shared.mem_mon.read()}') - log_debug(f'VAE decode: name={sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "baked"} dtype={model.vae.dtype} upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} images={latents.shape[0]} latents={latents.shape} time={round(t1-t0, 3)}') + shared.log.debug(f'VAE decode: {stats} time={round(t1-t0, 3)}') return decoded