90e911fd54
Only build the parser once (at import time) instead of for each step. doctest is run by simply executing modules/prompt_parser.py
251 lines
8.1 KiB
Python
251 lines
8.1 KiB
Python
import re
|
|
from collections import namedtuple
|
|
|
|
import lark
|
|
|
|
# a prompt like this: "fantasy landscape with a [mountain:lake:0.25] and [an oak:a christmas tree:0.75][ in foreground::0.6][ in background:0.25] [shoddy:masterful:0.5]"
|
|
# will be represented with prompt_schedule like this (assuming steps=100):
|
|
# [25, 'fantasy landscape with a mountain and an oak in foreground shoddy']
|
|
# [50, 'fantasy landscape with a lake and an oak in foreground in background shoddy']
|
|
# [60, 'fantasy landscape with a lake and an oak in foreground in background masterful']
|
|
# [75, 'fantasy landscape with a lake and an oak in background masterful']
|
|
# [100, 'fantasy landscape with a lake and a christmas tree in background masterful']
|
|
|
|
schedule_parser = lark.Lark(r"""
|
|
!start: (prompt | /[][():]/+)*
|
|
prompt: (emphasized | scheduled | plain | WHITESPACE)*
|
|
!emphasized: "(" prompt ")"
|
|
| "(" prompt ":" prompt ")"
|
|
| "[" prompt "]"
|
|
scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER "]"
|
|
WHITESPACE: /\s+/
|
|
plain: /([^\\\[\]():]|\\.)+/
|
|
%import common.SIGNED_NUMBER -> NUMBER
|
|
""")
|
|
|
|
def get_learned_conditioning_prompt_schedules(prompts, steps):
|
|
"""
|
|
>>> g = lambda p: get_learned_conditioning_prompt_schedules([p], 10)[0]
|
|
>>> g("test")
|
|
[[10, 'test']]
|
|
>>> g("a [b:3]")
|
|
[[3, 'a '], [10, 'a b']]
|
|
>>> g("a [b: 3]")
|
|
[[3, 'a '], [10, 'a b']]
|
|
>>> g("a [[[b]]:2]")
|
|
[[2, 'a '], [10, 'a [[b]]']]
|
|
>>> g("[(a:2):3]")
|
|
[[3, ''], [10, '(a:2)']]
|
|
>>> g("a [b : c : 1] d")
|
|
[[1, 'a b d'], [10, 'a c d']]
|
|
>>> g("a[b:[c:d:2]:1]e")
|
|
[[1, 'abe'], [2, 'ace'], [10, 'ade']]
|
|
>>> g("a [unbalanced")
|
|
[[10, 'a [unbalanced']]
|
|
>>> g("a [b:.5] c")
|
|
[[5, 'a c'], [10, 'a b c']]
|
|
>>> g("a [{b|d{:.5] c") # not handling this right now
|
|
[[5, 'a c'], [10, 'a {b|d{ c']]
|
|
>>> g("((a][:b:c [d:3]")
|
|
[[3, '((a][:b:c '], [10, '((a][:b:c d']]
|
|
"""
|
|
|
|
def collect_steps(steps, tree):
|
|
l = [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]))
|
|
l.append(tree.children[-1])
|
|
CollectSteps().visit(tree)
|
|
return sorted(set(l))
|
|
|
|
def at_step(step, tree):
|
|
class AtStep(lark.Transformer):
|
|
def scheduled(self, args):
|
|
before, after, _, when = args
|
|
yield before or () if step <= when else after
|
|
def start(self, args):
|
|
def flatten(x):
|
|
if type(x) == str:
|
|
yield x
|
|
else:
|
|
for gen in x:
|
|
yield from flatten(gen)
|
|
return ''.join(flatten(args))
|
|
def plain(self, args):
|
|
yield args[0].value
|
|
def __default__(self, data, children, meta):
|
|
for child in children:
|
|
yield from child
|
|
return AtStep().transform(tree)
|
|
|
|
def get_schedule(prompt):
|
|
try:
|
|
tree = schedule_parser.parse(prompt)
|
|
except lark.exceptions.LarkError as e:
|
|
if 0:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return [[steps, prompt]]
|
|
return [[t, at_step(t, tree)] for t in collect_steps(steps, tree)]
|
|
|
|
promptdict = {prompt: get_schedule(prompt) for prompt in set(prompts)}
|
|
return [promptdict[prompt] for prompt in prompts]
|
|
|
|
|
|
ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])
|
|
ScheduledPromptBatch = namedtuple("ScheduledPromptBatch", ["shape", "schedules"])
|
|
|
|
|
|
def get_learned_conditioning(model, prompts, steps):
|
|
res = []
|
|
|
|
prompt_schedules = get_learned_conditioning_prompt_schedules(prompts, steps)
|
|
cache = {}
|
|
|
|
for prompt, prompt_schedule in zip(prompts, prompt_schedules):
|
|
|
|
cached = cache.get(prompt, None)
|
|
if cached is not None:
|
|
res.append(cached)
|
|
continue
|
|
|
|
texts = [x[1] for x in prompt_schedule]
|
|
conds = model.get_learned_conditioning(texts)
|
|
|
|
cond_schedule = []
|
|
for i, (end_at_step, text) in enumerate(prompt_schedule):
|
|
cond_schedule.append(ScheduledPromptConditioning(end_at_step, conds[i]))
|
|
|
|
cache[prompt] = cond_schedule
|
|
res.append(cond_schedule)
|
|
|
|
return ScheduledPromptBatch((len(prompts),) + res[0][0].cond.shape, res)
|
|
|
|
|
|
def reconstruct_cond_batch(c: ScheduledPromptBatch, current_step):
|
|
param = c.schedules[0][0].cond
|
|
res = torch.zeros(c.shape, device=param.device, dtype=param.dtype)
|
|
for i, cond_schedule in enumerate(c.schedules):
|
|
target_index = 0
|
|
for current, (end_at, cond) in enumerate(cond_schedule):
|
|
if current_step <= end_at:
|
|
target_index = current
|
|
break
|
|
res[i] = cond_schedule[target_index].cond
|
|
|
|
return res
|
|
|
|
|
|
re_attention = re.compile(r"""
|
|
\\\(|
|
|
\\\)|
|
|
\\\[|
|
|
\\]|
|
|
\\\\|
|
|
\\|
|
|
\(|
|
|
\[|
|
|
:([+-]?[.\d]+)\)|
|
|
\)|
|
|
]|
|
|
[^\\()\[\]:]+|
|
|
:
|
|
""", re.X)
|
|
|
|
|
|
def parse_prompt_attention(text):
|
|
"""
|
|
Parses a string with attention tokens and returns a list of pairs: text and its assoicated weight.
|
|
Accepted tokens are:
|
|
(abc) - increases attention to abc by a multiplier of 1.1
|
|
(abc:3.12) - increases attention to abc by a multiplier of 3.12
|
|
[abc] - decreases attention to abc by a multiplier of 1.1
|
|
\( - literal character '('
|
|
\[ - literal character '['
|
|
\) - literal character ')'
|
|
\] - literal character ']'
|
|
\\ - literal character '\'
|
|
anything else - just text
|
|
|
|
>>> parse_prompt_attention('normal text')
|
|
[['normal text', 1.0]]
|
|
>>> parse_prompt_attention('an (important) word')
|
|
[['an ', 1.0], ['important', 1.1], [' word', 1.0]]
|
|
>>> parse_prompt_attention('(unbalanced')
|
|
[['unbalanced', 1.1]]
|
|
>>> parse_prompt_attention('\(literal\]')
|
|
[['(literal]', 1.0]]
|
|
>>> parse_prompt_attention('(unnecessary)(parens)')
|
|
[['unnecessaryparens', 1.1]]
|
|
>>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
|
|
[['a ', 1.0],
|
|
['house', 1.5730000000000004],
|
|
[' ', 1.1],
|
|
['on', 1.0],
|
|
[' a ', 1.1],
|
|
['hill', 0.55],
|
|
[', sun, ', 1.1],
|
|
['sky', 1.4641000000000006],
|
|
['.', 1.1]]
|
|
"""
|
|
|
|
res = []
|
|
round_brackets = []
|
|
square_brackets = []
|
|
|
|
round_bracket_multiplier = 1.1
|
|
square_bracket_multiplier = 1 / 1.1
|
|
|
|
def multiply_range(start_position, multiplier):
|
|
for p in range(start_position, len(res)):
|
|
res[p][1] *= multiplier
|
|
|
|
for m in re_attention.finditer(text):
|
|
text = m.group(0)
|
|
weight = m.group(1)
|
|
|
|
if text.startswith('\\'):
|
|
res.append([text[1:], 1.0])
|
|
elif text == '(':
|
|
round_brackets.append(len(res))
|
|
elif text == '[':
|
|
square_brackets.append(len(res))
|
|
elif weight is not None and len(round_brackets) > 0:
|
|
multiply_range(round_brackets.pop(), float(weight))
|
|
elif text == ')' and len(round_brackets) > 0:
|
|
multiply_range(round_brackets.pop(), round_bracket_multiplier)
|
|
elif text == ']' and len(square_brackets) > 0:
|
|
multiply_range(square_brackets.pop(), square_bracket_multiplier)
|
|
else:
|
|
res.append([text, 1.0])
|
|
|
|
for pos in round_brackets:
|
|
multiply_range(pos, round_bracket_multiplier)
|
|
|
|
for pos in square_brackets:
|
|
multiply_range(pos, square_bracket_multiplier)
|
|
|
|
if len(res) == 0:
|
|
res = [["", 1.0]]
|
|
|
|
# merge runs of identical weights
|
|
i = 0
|
|
while i + 1 < len(res):
|
|
if res[i][1] == res[i + 1][1]:
|
|
res[i][0] += res[i + 1][0]
|
|
res.pop(i + 1)
|
|
else:
|
|
i += 1
|
|
|
|
return res
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
|
|
else:
|
|
import torch # doctest faster
|