always keep lora on gpu

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  

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)

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
 
 

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
 
 

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
 

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

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)
 

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'

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