add support for whitespace after the number in constructions like [foo:bar: 0.5 ] and (foo : 0.5 )

1 files changed, 8 insertions, 8 deletions

diff --git a/modules/prompt_parser.py b/modules/prompt_parser.py
index 203ae1ac..8169a459 100644
--- a/modules/prompt_parser.py
+++ b/modules/prompt_parser.py
@@ -19,7 +19,7 @@ prompt: (emphasized | scheduled | alternate | plain | WHITESPACE)*
 !emphasized: "(" prompt ")"

         | "(" prompt ":" prompt ")"

         | "[" prompt "]"

-scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER "]"

+scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER [WHITESPACE] "]"

 alternate: "[" prompt ("|" prompt)+ "]"

 WHITESPACE: /\s+/

 plain: /([^\\\[\]():|]|\\.)+/

@@ -60,11 +60,11 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
 

         class CollectSteps(lark.Visitor):

             def scheduled(self, tree):

-                tree.children[-1] = float(tree.children[-1])

-                if tree.children[-1] < 1:

-                    tree.children[-1] *= steps

-                tree.children[-1] = min(steps, int(tree.children[-1]))

-                res.append(tree.children[-1])

+                tree.children[-2] = float(tree.children[-2])

+                if tree.children[-2] < 1:

+                    tree.children[-2] *= steps

+                tree.children[-2] = min(steps, int(tree.children[-2]))

+                res.append(tree.children[-2])

 

             def alternate(self, tree):

                 res.extend(range(1, steps+1))

@@ -75,7 +75,7 @@ def get_learned_conditioning_prompt_schedules(prompts, steps):
     def at_step(step, tree):

         class AtStep(lark.Transformer):

             def scheduled(self, args):

-                before, after, _, when = args

+                before, after, _, when, _ = args

                 yield before or () if step <= when else after

             def alternate(self, args):

                 yield next(args[(step - 1)%len(args)])

@@ -333,7 +333,7 @@ re_attention = re.compile(r"""
 \\|

 \(|

 \[|

-:([+-]?[.\d]+)\)|

+:\s*([+-]?[.\d]+)\s*\)|

 \)|

 ]|

 [^\\()\[\]:]+|