Fix processing with A1111 hires fix. (#4)

# Pull Request

## Description

Fixes processing A1111 hr prompts.

Fixes #3 

## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)

## How Has This Been Tested?

**Test Configuration**:

- A1111 v1.6
- SD.Next

## Checklist

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged and published in downstream
modules
- [x] I have checked my code and corrected any misspellings

.vscode/settings.json

@@ -1,4 +1,5 @@
 {
+    "python.analysis.extraPaths": ["../.."],
     "python.testing.unittestArgs": [
         "-v",
         "-s",
@@ -9,8 +10,6 @@
     "python.testing.pytestEnabled": false,
     "python.testing.unittestEnabled": true,
     "python.analysis.typeCheckingMode": "basic",
-    "python.linting.pylintEnabled": true,
-    "python.linting.enabled": true,
     "black-formatter.args": [
         "--line-length=120"
     ]

README.md

@@ -21,6 +21,16 @@ you will have to rename one of the folders, so that it loads before than "Send t
 
 When in doubt, just rename this extension's folder with a "z" in front (for example) so that it is the last one to load, or manually set such folder name when installing it.
 
+Note: it only recognizes regular A1111 prompt formats. So:
+
+* Attention: \[prompt\] (prompt) (prompt:weight)
+* Alternation: \[prompt1|prompt2|...\]
+* Scheduling: \[prompt1:prompt2:step\]
+
+In SD.Next that means only the A1111 or Full parsers.
+
+It does not build AND/BREAK separations into the negative prompt.
+
 ## Installation
 
 1. Go to Extensions > Install from URL

scripts/extension.py

@@ -13,26 +13,35 @@ from modules import scripts, shared, script_callbacks
 from modules.processing import StableDiffusionProcessing
 from modules.shared import opts
 from sendtonegative import SendToNegative
+from stnlogging import SendToNegativeLogFactory
 
 
 class SendToNegativeScript(scripts.Script):
     def __init__(self):
         if not hasattr(self, "callbacks_added"):
+            lf = SendToNegativeLogFactory()
+            self.__logstn = lf.log
             script_callbacks.on_ui_settings(self.__on_ui_settings)
             self.callbacks_added = True
 
     def title(self):
-        return f"{SendToNegative.NAME} v{SendToNegative.VERSION}"
+        return SendToNegative.NAME
 
     def show(self, is_img2img):
         return scripts.AlwaysVisible
 
     def process(self, p: StableDiffusionProcessing, *args, **kwargs):
-        stn = SendToNegative(opts=opts)
+        stn = SendToNegative(self.__logstn, opts=opts)
         for i in range(len(p.all_prompts)):  # pylint: disable=consider-using-enumerate
             p.all_prompts[i], p.all_negative_prompts[i] = stn.process_prompt(
                 p.all_prompts[i], p.all_negative_prompts[i]
             )
+        # make it compatible with A1111 hires fix
+        if hasattr(p, "all_hr_prompts") and hasattr(p, "all_hr_negative_prompts"):
+            for i in range(len(p.all_hr_prompts)):  # pylint: disable=consider-using-enumerate
+                p.all_hr_prompts[i], p.all_hr_negative_prompts[i] = stn.process_prompt(
+                    p.all_hr_prompts[i], p.all_hr_negative_prompts[i]
+                )
 
     def __on_ui_settings(self):
         section = ("send-to-negative", SendToNegative.NAME)

sendtonegative.py

@@ -1,17 +1,18 @@
 from collections import namedtuple
-import logging
-import lark
 import re
+import math
+import lark
 
 
 class SendToNegative:  # pylint: disable=too-few-public-methods
     NAME = "Send to Negative"
-    VERSION = "2.0"
+    VERSION = "2.1.4"
 
     DEFAULT_SEPARATOR = ", "
 
     def __init__(
         self,
+        log,
         separator=None,
         ignore_repeats=None,
         join_attention=None,
@@ -33,7 +34,9 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
 
             iN - tags the position of insertion point N. Used only in the negative prompt and does not accept content. N can be 0 to 9.
         """
-        self.__logger = logging.getLogger(__name__)
+        self.__logger = log
+        if opts is not None and getattr(opts, "prompt_attention", "") == "Compel parser":
+            self.__logger.warning("Compel parser is not supported!")
         self.__ignore_repeats = (
             ignore_repeats if ignore_repeats is not None else getattr(opts, "stn_ignorerepeats", True)
         )
@@ -94,6 +97,9 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
             self.add_at = add_at
             self.remove = []
 
+        def __get_numpar_value(self, numpar):
+            return float(next(x for x in numpar.children if x.type == "NUMBER").value)
+
         def scheduled(self, tree):
             if len(tree.children) > 2:  # before & after
                 before = tree.children[0]
@@ -101,20 +107,20 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
                 before = None
             after = tree.children[-2]
             numpar = tree.children[-1]
-            pos = float(numpar.children[0].value)
+            pos = self.__get_numpar_value(numpar)
             if pos >= 1:
                 pos = int(pos)
             # self.__shell.append(self.AccumulatedShell("sc", tree.meta.start_pos, pos))
             if before is not None and hasattr(before, "data"):
                 self.__logger.debug(
-                    f"Shell scheduled before at {[before.meta.start_pos,before.meta.end_pos] if hasattr(before,'meta') else '?'} : {pos}"
+                    f"Shell scheduled before at {[before.meta.start_pos,before.meta.end_pos] if hasattr(before,'meta') and not before.meta.empty else '?'} : {pos}"
                 )
                 self.__shell.append(self.AccumulatedShell("scb", pos, None))
                 self.visit(before)
                 self.__shell.pop()
             if hasattr(after, "data"):
                 self.__logger.debug(
-                    f"Shell scheduled after at {[after.meta.start_pos,after.meta.end_pos] if hasattr(after,'meta') else '?'} : {pos}"
+                    f"Shell scheduled after at {[after.meta.start_pos,after.meta.end_pos] if hasattr(after,'meta') and not after.meta.empty else '?'} : {pos}"
                 )
                 self.__shell.append(self.AccumulatedShell("sca", pos, None))
                 self.visit(after)
@@ -125,7 +131,7 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
             # self.__shell.append(self.AccumulatedShell("al", tree.meta.start_pos, len(tree.children)))
             for i, opt in enumerate(tree.children):
                 self.__logger.debug(
-                    f"Shell alternate at {[opt.meta.start_pos,opt.meta.end_pos] if hasattr(opt,'meta') else '?'} : {i+1}"
+                    f"Shell alternate at {[opt.meta.start_pos,opt.meta.end_pos] if hasattr(opt,'meta') and not opt.meta.empty else '?'} : {i+1}"
                 )
                 if hasattr(opt, "data"):
                     self.__shell.append(self.AccumulatedShell("alo", i + 1, len(tree.children)))
@@ -135,9 +141,9 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
 
         def emphasized(self, tree):
             numpar = tree.children[-1]
-            weight = float(numpar.children[0].value) if numpar is not None else 1.1
+            weight = self.__get_numpar_value(numpar) if numpar is not None else 1.1
             self.__logger.debug(
-                f"Shell attention at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') else '?'}: {weight}"
+                f"Shell attention at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') and not tree.meta.empty else '?'}: {weight}"
             )
             self.__shell.append(self.AccumulatedShell("at", weight, None))
             self.visit_children(tree)
@@ -146,7 +152,7 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
         def deemphasized(self, tree):
             weight = 0.9
             self.__logger.debug(
-                f"Shell attention at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') else '?'}: {weight}"
+                f"Shell attention at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') and not tree.meta.empty else '?'}: {weight}"
             )
             self.__shell.append(self.AccumulatedShell("at", weight, None))
             self.visit_children(tree)
@@ -157,13 +163,17 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
             parameters = negtagparameters.children[0].value if negtagparameters is not None else ""
             rest = []
             for x in tree.children[1::]:
-                rest.append(self.__prompt[x.meta.start_pos : x.meta.end_pos] if hasattr(x, "meta") else x.value)
+                rest.append(
+                    self.__prompt[x.meta.start_pos : x.meta.end_pos]
+                    if hasattr(x, "meta") and not x.meta.empty
+                    else x.value
+                )
             content = "".join(rest)
             self.__negtags.append(
                 self.NegTag(tree.meta.start_pos, tree.meta.end_pos, content, parameters, self.__shell.copy())
             )
             self.__logger.debug(
-                f"Negative tag at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') else '?'}: {parameters}: {content.encode('unicode_escape').decode('utf-8')}"
+                f"Negative tag at {[tree.meta.start_pos,tree.meta.end_pos] if hasattr(tree,'meta') and not tree.meta.empty else '?'}: {parameters}: {content.encode('unicode_escape').decode('utf-8')}"
             )
 
         def start(self, tree):
@@ -176,7 +186,8 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
                         if nt.shell[i].type == "at" and nt.shell[i - 1].type == "at":
                             nt.shell[i - 1] = self.AccumulatedShell(
                                 "at",
-                                (100 * nt.shell[i - 1].info1 * nt.shell[i].info1) / 100,  # we limit to two decimals
+                                math.floor(100 * nt.shell[i - 1].info1 * nt.shell[i].info1)
+                                / 100,  # we limit to two decimals
                                 None,
                             )
                             nt.shell.pop(i)
@@ -194,14 +205,14 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
                             else:
                                 start += "("
                                 end = f":{s.info1})" + end
-                        #case "sc":
+                        # case "sc":
                         case "scb":
                             start += "["
                             end = f"::{s.info1}]" + end
                         case "sca":
                             start += "["
                             end = f":{s.info1}]" + end
-                        #case "al":
+                        # case "al":
                         case "alo":
                             start += "[" + ("|" * int(s.info1 - 1))
                             end = ("|" * int(s.info2 - s.info1)) + "]" + end
@@ -253,7 +264,7 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
     def __find_tags(self, prompt):
         add_at = {"start": [], "insertion_point": [[] for x in range(10)], "end": []}
         tree = self.__schedule_parser.parse(prompt)
-        self.__logger.debug(f"Initial tree: {tree.pretty()}")
+        self.__logger.debug(f"Initial tree:\n{tree.pretty()}")
 
         readtree = self.ReadTree(self.__logger, self.__ignore_repeats, self.__join_attention, prompt, add_at)
         readtree.visit(tree)
@@ -261,9 +272,9 @@ class SendToNegative:  # pylint: disable=too-few-public-methods
         for r in readtree.remove[::-1]:
             prompt = prompt[: r[0]] + prompt[r[1] :]
         if self.__cleanup:
-            prompt = re.sub(r"\((?::[\d\.]+)?\)", "", prompt)  # clean up empty attention
+            prompt = re.sub(r"\((?::[+-]?[\d\.]+)?\)", "", prompt)  # clean up empty attention
             prompt = re.sub(r"\[\]", "", prompt)  # clean up empty attention
-            prompt = re.sub(r"\[:?:[\d\.]+\]", "", prompt)  # clean up empty scheduling
+            prompt = re.sub(r"\[:?:[+-]?[\d\.]+\]", "", prompt)  # clean up empty scheduling
             prompt = re.sub(r"\[\|+\]", "", prompt)  # clean up empty alternation
             # clean up whitespace and extra separators
             prompt = (

stnlogging.py

@@ -0,0 +1,40 @@
+import sys
+import copy
+import logging
+
+
+class SendToNegativeLogFactory:
+    class ColoredFormatter(logging.Formatter):
+        COLORS = {
+            "DEBUG": "\033[0;36m",  # CYAN
+            "INFO": "\033[0;32m",  # GREEN
+            "WARNING": "\033[0;33m",  # YELLOW
+            "ERROR": "\033[0;31m",  # RED
+            "CRITICAL": "\033[0;37;41m",  # WHITE ON RED
+            "RESET": "\033[0m",  # RESET COLOR
+        }
+
+        def format(self, record):
+            colored_record = copy.copy(record)
+            levelname = colored_record.levelname
+            seq = self.COLORS.get(levelname, self.COLORS["RESET"])
+            colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
+            return super().format(colored_record)
+
+    def __init__(self):
+        logsd = logging.getLogger("sd")
+        stnlog = logging.getLogger("SendToNegative")
+        stnlog.setLevel(logging.INFO)
+        stnlog.handlers = logsd.handlers
+        if not stnlog.handlers:
+            handler = logging.StreamHandler(sys.stdout)
+            handler.setFormatter(self.ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+            stnlog.addHandler(handler)
+            self.log = stnlog
+        else:
+            self.log = SendToNegativeLogCustomAdapter(stnlog)
+
+
+class SendToNegativeLogCustomAdapter(logging.LoggerAdapter):
+    def process(self, msg, kwargs):
+        return f"[SendToNegative] {msg}", kwargs

tests/tests.py

@@ -6,17 +6,15 @@ import os
 sys.path.insert(1, os.path.join(sys.path[0], ".."))
 
 from sendtonegative import SendToNegative  # pylint: disable=import-error
+from stnlogging import SendToNegativeLogFactory
 
 
 class TestSendToNegative(unittest.TestCase):
     def setUp(self):
-        self.defstn = SendToNegative(
-            separator=", ",
-            ignore_repeats=True,
-            join_attention=True,
-            cleanup=True,
-        )
-        logging.basicConfig(level=logging.DEBUG)
+        lf = SendToNegativeLogFactory()
+        self.__log = lf.log
+        self.__log.setLevel(logging.DEBUG)
+        self.defstn = SendToNegative(self.__log, separator=", ", ignore_repeats=True, join_attention=True, cleanup=True)
 
     def process(
         self,
@@ -121,19 +119,14 @@ class TestSendToNegative(unittest.TestCase):
             "normal quality, <!!i0!!>, bad quality<!!i1!!>, worse quality",
             " (), flowers , ,  ",
             "red, (pink), normal quality, yellow, bad quality, green, worse quality, purple, blue",
-            SendToNegative(
-                separator=", ",
-                ignore_repeats=True,
-                join_attention=True,
-                cleanup=False,
-            ),
+            SendToNegative(self.__log, separator=", ", ignore_repeats=True, join_attention=True, cleanup=False),
         )
 
     def test_inside_attention1(self):
         self.process(
-            "[<!neg1!>] this is a ((test<!!e!neg2!>) (test:2.0):1.5)",
+            "[<!neg1!>] this is a ((test<!!e!neg2!>) (test:2.0): 1.5 )",
             "normal quality",
-            "this is a ((test) (test:2.0):1.5)",
+            "this is a ((test) (test:2.0): 1.5 )",
             "[neg1], normal quality, (neg2:1.65)",
         )
 
@@ -163,17 +156,17 @@ class TestSendToNegative(unittest.TestCase):
 
     def test_inside_alternation3(self):
         self.process(
-            "this is a (([complex<!neg1!>[one|two<!neg12!>|three|four(<!neg14!>)]|simple<!neg2!>|regular<!neg3!>] test)(test:2.0):1.5)",
+            "this is a (([complex<!neg1!>[one|two<!neg12!>||three|four(<!neg14!>)]|simple<!neg2!>|regular<!neg3!>] test)(test:2.0):1.5)",
             "normal quality",
-            "this is a (([complex[one|two|three|four]|simple|regular] test)(test:2.0):1.5)",
-            "([neg1||]:1.65), ([[|neg12||]||]:1.65), ([[|||(neg14)]||]:1.65), ([|neg2|]:1.65), ([||neg3]:1.65), normal quality",
+            "this is a (([complex[one|two||three|four]|simple|regular] test)(test:2.0):1.5)",
+            "([neg1||]:1.65), ([[|neg12|||]||]:1.65), ([[||||(neg14)]||]:1.65), ([|neg2|]:1.65), ([||neg3]:1.65), normal quality",
         )
 
     def test_inside_scheduling(self):
         self.process(
-            "this is [abc<!neg1!>:def<!!e!neg2!>:5]",
+            "this is [abc<!neg1!>:def<!!e!neg2!>: 5 ]",
             "normal quality",
-            "this is [abc:def:5]",
+            "this is [abc:def: 5 ]",
             "[neg1::5], normal quality, [neg2:5]",
         )