7 files changed, 227 insertions, 133 deletions
diff --git a/extensions-builtin/LDSR/scripts/ldsr_model.py b/extensions-builtin/LDSR/scripts/ldsr_model.py
index b8cff29b..da19cff1 100644
--- a/extensions-builtin/LDSR/scripts/ldsr_model.py
+++ b/extensions-builtin/LDSR/scripts/ldsr_model.py
@@ -25,22 +25,28 @@ class UpscalerLDSR(Upscaler):
         yaml_path = os.path.join(self.model_path, "project.yaml")
         old_model_path = os.path.join(self.model_path, "model.pth")
         new_model_path = os.path.join(self.model_path, "model.ckpt")
-        safetensors_model_path = os.path.join(self.model_path, "model.safetensors")
+
+        local_model_paths = self.find_models(ext_filter=[".ckpt", ".safetensors"])
+        local_ckpt_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.ckpt")]), None)
+        local_safetensors_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.safetensors")]), None)
+        local_yaml_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("project.yaml")]), None)
+
         if os.path.exists(yaml_path):
             statinfo = os.stat(yaml_path)
             if statinfo.st_size >= 10485760:
                 print("Removing invalid LDSR YAML file.")
                 os.remove(yaml_path)
+
         if os.path.exists(old_model_path):
             print("Renaming model from model.pth to model.ckpt")
             os.rename(old_model_path, new_model_path)
-        if os.path.exists(safetensors_model_path):
-            model = safetensors_model_path
+
+        if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
+            model = local_safetensors_path
         else:
-            model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
-                                       file_name="model.ckpt", progress=True)
-        yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,
-                                  file_name="project.yaml", progress=True)
+            model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="model.ckpt", progress=True)
+
+        yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_path, file_name="project.yaml", progress=True)
 
         try:
             return LDSR(model, yaml)
diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py
index 6be6ef73..ccb249ac 100644
--- a/extensions-builtin/Lora/extra_networks_lora.py
+++ b/extensions-builtin/Lora/extra_networks_lora.py
@@ -1,6 +1,7 @@
 from modules import extra_networks, shared
 import lora
 
+
 class ExtraNetworkLora(extra_networks.ExtraNetwork):
     def __init__(self):
         super().__init__('lora')
@@ -8,7 +9,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
     def activate(self, p, params_list):
         additional = shared.opts.sd_lora
 
-        if additional != "" and additional in lora.available_loras and len([x for x in params_list if x.items[0] == additional]) == 0:
+        if additional != "None" and additional in lora.available_loras and len([x for x in params_list if x.items[0] == additional]) == 0:
             p.all_prompts = [x + f"<lora:{additional}:{shared.opts.extra_networks_default_multiplier}>" for x in p.all_prompts]
             params_list.append(extra_networks.ExtraNetworkParams(items=[additional, shared.opts.extra_networks_default_multiplier]))
 
diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py
index d3eb0d3b..d488b5ae 100644
--- a/extensions-builtin/Lora/lora.py
+++ b/extensions-builtin/Lora/lora.py
@@ -4,7 +4,7 @@ import re
 import torch
 from typing import Union
 
-from modules import shared, devices, sd_models, errors
+from modules import shared, devices, sd_models, errors, scripts
 
 metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
 
@@ -93,6 +93,7 @@ class LoraOnDisk:
             self.metadata = m
 
         self.ssmd_cover_images = self.metadata.pop('ssmd_cover_images', None)  # those are cover images and they are too big to display in UI as text
+        self.alias = self.metadata.get('ss_output_name', self.name)
 
 
 class LoraModule:
@@ -165,8 +166,10 @@ def load_lora(name, filename):
             module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
         elif type(sd_module) == torch.nn.MultiheadAttention:
             module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
-        elif type(sd_module) == torch.nn.Conv2d:
+        elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (1, 1):
             module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
+        elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (3, 3):
+            module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (3, 3), bias=False)
         else:
             print(f'Lora layer {key_diffusers} matched a layer with unsupported type: {type(sd_module).__name__}')
             continue
@@ -199,11 +202,11 @@ def load_loras(names, multipliers=None):
 
     loaded_loras.clear()
 
-    loras_on_disk = [available_loras.get(name, None) for name in names]
+    loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
     if any([x is None for x in loras_on_disk]):
         list_available_loras()
 
-        loras_on_disk = [available_loras.get(name, None) for name in names]
+        loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
 
     for i, name in enumerate(names):
         lora = already_loaded.get(name, None)
@@ -211,7 +214,11 @@ def load_loras(names, multipliers=None):
         lora_on_disk = loras_on_disk[i]
         if lora_on_disk is not None:
             if lora is None or os.path.getmtime(lora_on_disk.filename) > lora.mtime:
-                lora = load_lora(name, lora_on_disk.filename)
+                try:
+                    lora = load_lora(name, lora_on_disk.filename)
+                except Exception as e:
+                    errors.display(e, f"loading Lora {lora_on_disk.filename}")
+                    continue
 
         if lora is None:
             print(f"Couldn't find Lora with name {name}")
@@ -228,6 +235,8 @@ def lora_calc_updown(lora, module, target):
 
         if up.shape[2:] == (1, 1) and down.shape[2:] == (1, 1):
             updown = (up.squeeze(2).squeeze(2) @ down.squeeze(2).squeeze(2)).unsqueeze(2).unsqueeze(3)
+        elif up.shape[2:] == (3, 3) or down.shape[2:] == (3, 3):
+            updown = torch.nn.functional.conv2d(down.permute(1, 0, 2, 3), up).permute(1, 0, 2, 3)
         else:
             updown = up @ down
 
@@ -236,6 +245,19 @@ def lora_calc_updown(lora, module, target):
         return updown
 
 
+def lora_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
+    weights_backup = getattr(self, "lora_weights_backup", None)
+
+    if weights_backup is None:
+        return
+
+    if isinstance(self, torch.nn.MultiheadAttention):
+        self.in_proj_weight.copy_(weights_backup[0])
+        self.out_proj.weight.copy_(weights_backup[1])
+    else:
+        self.weight.copy_(weights_backup)
+
+
 def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
     """
     Applies the currently selected set of Loras to the weights of torch layer self.
@@ -260,12 +282,7 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu
         self.lora_weights_backup = weights_backup
 
     if current_names != wanted_names:
-        if weights_backup is not None:
-            if isinstance(self, torch.nn.MultiheadAttention):
-                self.in_proj_weight.copy_(weights_backup[0])
-                self.out_proj.weight.copy_(weights_backup[1])
-            else:
-                self.weight.copy_(weights_backup)
+        lora_restore_weights_from_backup(self)
 
         for lora in loaded_loras:
             module = lora.modules.get(lora_layer_name, None)
@@ -296,12 +313,45 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu
         setattr(self, "lora_current_names", wanted_names)
 
 
+def lora_forward(module, input, original_forward):
+    """
+    Old way of applying Lora by executing operations during layer's forward.
+    Stacking many loras this way results in big performance degradation.
+    """
+
+    if len(loaded_loras) == 0:
+        return original_forward(module, input)
+
+    input = devices.cond_cast_unet(input)
+
+    lora_restore_weights_from_backup(module)
+    lora_reset_cached_weight(module)
+
+    res = original_forward(module, input)
+
+    lora_layer_name = getattr(module, 'lora_layer_name', None)
+    for lora in loaded_loras:
+        module = lora.modules.get(lora_layer_name, None)
+        if module is None:
+            continue
+
+        module.up.to(device=devices.device)
+        module.down.to(device=devices.device)
+
+        res = res + module.up(module.down(input)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
+
+    return res
+
+
 def lora_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]):
     setattr(self, "lora_current_names", ())
     setattr(self, "lora_weights_backup", None)
 
 
 def lora_Linear_forward(self, input):
+    if shared.opts.lora_functional:
+        return lora_forward(self, input, torch.nn.Linear_forward_before_lora)
+
     lora_apply_weights(self)
 
     return torch.nn.Linear_forward_before_lora(self, input)
@@ -314,6 +364,9 @@ def lora_Linear_load_state_dict(self, *args, **kwargs):
 
 
 def lora_Conv2d_forward(self, input):
+    if shared.opts.lora_functional:
+        return lora_forward(self, input, torch.nn.Conv2d_forward_before_lora)
+
     lora_apply_weights(self)
 
     return torch.nn.Conv2d_forward_before_lora(self, input)
@@ -339,24 +392,60 @@ def lora_MultiheadAttention_load_state_dict(self, *args, **kwargs):
 
 def list_available_loras():
     available_loras.clear()
+    available_lora_aliases.clear()
 
     os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)
 
-    candidates = \
-        glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.pt'), recursive=True) + \
-        glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.safetensors'), recursive=True) + \
-        glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.ckpt'), recursive=True)
-
+    candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
     for filename in sorted(candidates, key=str.lower):
         if os.path.isdir(filename):
             continue
 
         name = os.path.splitext(os.path.basename(filename))[0]
+        entry = LoraOnDisk(name, filename)
+
+        available_loras[name] = entry
+
+        available_lora_aliases[name] = entry
+        available_lora_aliases[entry.alias] = entry
+
+
+re_lora_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
+
+
+def infotext_pasted(infotext, params):
+    if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]:
+        return  # if the other extension is active, it will handle those fields, no need to do anything
+
+    added = []
+
+    for k, v in params.items():
+        if not k.startswith("AddNet Model "):
+            continue
+
+        num = k[13:]
+
+        if params.get("AddNet Module " + num) != "LoRA":
+            continue
+
+        name = params.get("AddNet Model " + num)
+        if name is None:
+            continue
+
+        m = re_lora_name.match(name)
+        if m:
+            name = m.group(1)
+
+        multiplier = params.get("AddNet Weight A " + num, "1.0")
+
+        added.append(f"<lora:{name}:{multiplier}>")
 
-        available_loras[name] = LoraOnDisk(name, filename)
+    if added:
+        params["Prompt"] += "\n" + "".join(added)
 
 
 available_loras = {}
+available_lora_aliases = {}
 loaded_loras = []
 
 list_available_loras()
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index 0adab225..a67b8a69 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -49,8 +49,14 @@ torch.nn.MultiheadAttention._load_from_state_dict = lora.lora_MultiheadAttention
 script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
 script_callbacks.on_script_unloaded(unload)
 script_callbacks.on_before_ui(before_ui)
+script_callbacks.on_infotext_pasted(lora.infotext_pasted)
 
 
 shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), {
-    "sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": [""] + [x for x in lora.available_loras]}, refresh=lora.list_available_loras),
+    "sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in lora.available_loras]}, refresh=lora.list_available_loras),
+}))
+
+
+shared.options_templates.update(shared.options_section(('compatibility', "Compatibility"), {
+    "lora_functional": shared.OptionInfo(False, "Lora: use old method that takes longer when you have multiple Loras active and produces same results as kohya-ss/sd-webui-additional-networks extension"),
 }))
diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py
index 68b11332..a0edbc1e 100644
--- a/extensions-builtin/Lora/ui_extra_networks_lora.py
+++ b/extensions-builtin/Lora/ui_extra_networks_lora.py
@@ -21,7 +21,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
                 "preview": self.find_preview(path),
                 "description": self.find_description(path),
                 "search_term": self.search_terms_from_path(lora_on_disk.filename),
-                "prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
+                "prompt": json.dumps(f"<lora:{lora_on_disk.alias}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
                 "local_preview": f"{path}.{shared.opts.samples_format}",
                 "metadata": json.dumps(lora_on_disk.metadata, indent=4) if lora_on_disk.metadata else None,
             }
diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/extensions-builtin/ScuNET/scripts/scunet_model.py
index e0fbf3a3..c7fd5739 100644
--- a/extensions-builtin/ScuNET/scripts/scunet_model.py
+++ b/extensions-builtin/ScuNET/scripts/scunet_model.py
@@ -5,11 +5,15 @@ import traceback
 import PIL.Image
 import numpy as np
 import torch
+from tqdm import tqdm
+
 from basicsr.utils.download_util import load_file_from_url
 
 import modules.upscaler
 from modules import devices, modelloader
 from scunet_model_arch import SCUNet as net
+from modules.shared import opts
+from modules import images
 
 
 class UpscalerScuNET(modules.upscaler.Upscaler):
@@ -42,28 +46,78 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
             scalers.append(scaler_data2)
         self.scalers = scalers
 
-    def do_upscale(self, img: PIL.Image, selected_file):
+    @staticmethod
+    @torch.no_grad()
+    def tiled_inference(img, model):
+        # test the image tile by tile
+        h, w = img.shape[2:]
+        tile = opts.SCUNET_tile
+        tile_overlap = opts.SCUNET_tile_overlap
+        if tile == 0:
+            return model(img)
+
+        device = devices.get_device_for('scunet')
+        assert tile % 8 == 0, "tile size should be a multiple of window_size"
+        sf = 1
+
+        stride = tile - tile_overlap
+        h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
+        w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
+        E = torch.zeros(1, 3, h * sf, w * sf, dtype=img.dtype, device=device)
+        W = torch.zeros_like(E, dtype=devices.dtype, device=device)
+
+        with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="ScuNET tiles") as pbar:
+            for h_idx in h_idx_list:
+
+                for w_idx in w_idx_list:
+
+                    in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
+
+                    out_patch = model(in_patch)
+                    out_patch_mask = torch.ones_like(out_patch)
+
+                    E[
+                        ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
+                    ].add_(out_patch)
+                    W[
+                        ..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
+                    ].add_(out_patch_mask)
+                    pbar.update(1)
+        output = E.div_(W)
+
+        return output
+
+    def do_upscale(self, img: PIL.Image.Image, selected_file):
+
         torch.cuda.empty_cache()
 
         model = self.load_model(selected_file)
         if model is None:
+            print(f"ScuNET: Unable to load model from {selected_file}", file=sys.stderr)
             return img
 
         device = devices.get_device_for('scunet')
-        img = np.array(img)
-        img = img[:, :, ::-1]
-        img = np.moveaxis(img, 2, 0) / 255
-        img = torch.from_numpy(img).float()
-        img = img.unsqueeze(0).to(device)
-
-        with torch.no_grad():
-            output = model(img)
-        output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = 255. * np.moveaxis(output, 0, 2)
-        output = output.astype(np.uint8)
-        output = output[:, :, ::-1]
+        tile = opts.SCUNET_tile
+        h, w = img.height, img.width
+        np_img = np.array(img)
+        np_img = np_img[:, :, ::-1]  # RGB to BGR
+        np_img = np_img.transpose((2, 0, 1)) / 255  # HWC to CHW
+        torch_img = torch.from_numpy(np_img).float().unsqueeze(0).to(device)  # type: ignore
+
+        if tile > h or tile > w:
+            _img = torch.zeros(1, 3, max(h, tile), max(w, tile), dtype=torch_img.dtype, device=torch_img.device)
+            _img[:, :, :h, :w] = torch_img # pad image
+            torch_img = _img
+
+        torch_output = self.tiled_inference(torch_img, model).squeeze(0)
+        torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any
+        np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy()
+        del torch_img, torch_output
         torch.cuda.empty_cache()
-        return PIL.Image.fromarray(output, 'RGB')
+
+        output = np_output.transpose((1, 2, 0))  # CHW to HWC
+        output = output[:, :, ::-1]  # BGR to RGB
+        return PIL.Image.fromarray((output * 255).astype(np.uint8))
 
     def load_model(self, path: str):
         device = devices.get_device_for('scunet')
@@ -84,4 +138,3 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
         model = model.to(device)
 
         return model
-
diff --git a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
index f0918e26..5c7a836a 100644
--- a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
+++ b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
@@ -1,103 +1,42 @@
 // Stable Diffusion WebUI - Bracket checker
-// Version 1.0
-// By Hingashi no Florin/Bwin4L
+// By Hingashi no Florin/Bwin4L & @akx
 // Counts open and closed brackets (round, square, curly) in the prompt and negative prompt text boxes in the txt2img and img2img tabs.
 // If there's a mismatch, the keyword counter turns red and if you hover on it, a tooltip tells you what's wrong.
 
-function checkBrackets(evt, textArea, counterElt) {
-  errorStringParen  = '(...) - Different number of opening and closing parentheses detected.\n';
-  errorStringSquare = '[...] - Different number of opening and closing square brackets detected.\n';
-  errorStringCurly  = '{...} - Different number of opening and closing curly brackets detected.\n';
-
-  openBracketRegExp = /\(/g;
-  closeBracketRegExp = /\)/g;
-
-  openSquareBracketRegExp = /\[/g;
-  closeSquareBracketRegExp = /\]/g;
-
-  openCurlyBracketRegExp = /\{/g;
-  closeCurlyBracketRegExp = /\}/g;
-
-  totalOpenBracketMatches = 0;
-  totalCloseBracketMatches = 0;
-  totalOpenSquareBracketMatches = 0;
-  totalCloseSquareBracketMatches = 0;
-  totalOpenCurlyBracketMatches = 0;
-  totalCloseCurlyBracketMatches = 0;
-
-  openBracketMatches = textArea.value.match(openBracketRegExp);
-  if(openBracketMatches) {
-    totalOpenBracketMatches = openBracketMatches.length;
-  }
-
-  closeBracketMatches = textArea.value.match(closeBracketRegExp);
-  if(closeBracketMatches) {
-    totalCloseBracketMatches = closeBracketMatches.length;
-  }
-
-  openSquareBracketMatches = textArea.value.match(openSquareBracketRegExp);
-  if(openSquareBracketMatches) {
-    totalOpenSquareBracketMatches = openSquareBracketMatches.length;
-  }
-
-  closeSquareBracketMatches = textArea.value.match(closeSquareBracketRegExp);
-  if(closeSquareBracketMatches) {
-    totalCloseSquareBracketMatches = closeSquareBracketMatches.length;
-  }
-
-  openCurlyBracketMatches = textArea.value.match(openCurlyBracketRegExp);
-  if(openCurlyBracketMatches) {
-    totalOpenCurlyBracketMatches = openCurlyBracketMatches.length;
-  }
-
-  closeCurlyBracketMatches = textArea.value.match(closeCurlyBracketRegExp);
-  if(closeCurlyBracketMatches) {
-    totalCloseCurlyBracketMatches = closeCurlyBracketMatches.length;
-  }
-
-  if(totalOpenBracketMatches != totalCloseBracketMatches) {
-    if(!counterElt.title.includes(errorStringParen)) {
-      counterElt.title += errorStringParen;
-    }
-  } else {
-    counterElt.title = counterElt.title.replace(errorStringParen, '');
-  }
-
-  if(totalOpenSquareBracketMatches != totalCloseSquareBracketMatches) {
-    if(!counterElt.title.includes(errorStringSquare)) {
-      counterElt.title += errorStringSquare;
-    }
-  } else {
-    counterElt.title = counterElt.title.replace(errorStringSquare, '');
-  }
-
-  if(totalOpenCurlyBracketMatches != totalCloseCurlyBracketMatches) {
-    if(!counterElt.title.includes(errorStringCurly)) {
-      counterElt.title += errorStringCurly;
+function checkBrackets(textArea, counterElt) {
+  var counts = {};
+  (textArea.value.match(/[(){}\[\]]/g) || []).forEach(bracket => {
+    counts[bracket] = (counts[bracket] || 0) + 1;
+  });
+  var errors = [];
+
+  function checkPair(open, close, kind) {
+    if (counts[open] !== counts[close]) {
+      errors.push(
+        `${open}...${close} - Detected ${counts[open] || 0} opening and ${counts[close] || 0} closing ${kind}.`
+      );
     }
-  } else {
-    counterElt.title = counterElt.title.replace(errorStringCurly, '');
   }
 
-  if(counterElt.title != '') {
-    counterElt.classList.add('error');
-  } else {
-    counterElt.classList.remove('error');
-  }
+  checkPair('(', ')', 'round brackets');
+  checkPair('[', ']', 'square brackets');
+  checkPair('{', '}', 'curly brackets');
+  counterElt.title = errors.join('\n');
+  counterElt.classList.toggle('error', errors.length !== 0);
 }
 
-function setupBracketChecking(id_prompt, id_counter){
-    var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea");
-    var counter = gradioApp().getElementById(id_counter)
+function setupBracketChecking(id_prompt, id_counter) {
+  var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea");
+  var counter = gradioApp().getElementById(id_counter)
 
-    textarea.addEventListener("input", function(evt){
-        checkBrackets(evt, textarea, counter)
-    });
+  if (textarea && counter) {
+    textarea.addEventListener("input", () => checkBrackets(textarea, counter));
+  }
 }
 
-onUiLoaded(function(){
-    setupBracketChecking('txt2img_prompt', 'txt2img_token_counter')
-    setupBracketChecking('txt2img_neg_prompt', 'txt2img_negative_token_counter')
-    setupBracketChecking('img2img_prompt', 'img2img_token_counter')
-    setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter')
-})
-\ No newline at end of file
+onUiLoaded(function () {
+  setupBracketChecking('txt2img_prompt', 'txt2img_token_counter');
+  setupBracketChecking('txt2img_neg_prompt', 'txt2img_negative_token_counter');
+  setupBracketChecking('img2img_prompt', 'img2img_token_counter');
+  setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter');
+});