cymic/Waifu_Diffusion_Webui
1
1import re2from collections import namedtuple3from typing import List4import lark5 6# 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]"7# will be represented with prompt_schedule like this (assuming steps=100):8# [25, 'fantasy landscape with a mountain and an oak in foreground shoddy']9# [50, 'fantasy landscape with a lake and an oak in foreground in background shoddy']10# [60, 'fantasy landscape with a lake and an oak in foreground in background masterful']11# [75, 'fantasy landscape with a lake and an oak in background masterful']12# [100, 'fantasy landscape with a lake and a christmas tree in background masterful']13 14schedule_parser = lark.Lark(r"""15!start: (prompt | /[][():]/+)*16prompt: (emphasized | scheduled | plain | WHITESPACE)*17!emphasized: "(" prompt ")"18 | "(" prompt ":" prompt ")"19 | "[" prompt "]"20scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER "]"21WHITESPACE: /\s+/22plain: /([^\\\[\]():]|\\.)+/23%import common.SIGNED_NUMBER -> NUMBER24""")25 26def get_learned_conditioning_prompt_schedules(prompts, steps):27 """28 >>> g = lambda p: get_learned_conditioning_prompt_schedules([p], 10)[0]29 >>> g("test")30 [[10, 'test']]31 >>> g("a [b:3]")32 [[3, 'a '], [10, 'a b']]33 >>> g("a [b: 3]")34 [[3, 'a '], [10, 'a b']]35 >>> g("a [[[b]]:2]")36 [[2, 'a '], [10, 'a [[b]]']]37 >>> g("[(a:2):3]")38 [[3, ''], [10, '(a:2)']]39 >>> g("a [b : c : 1] d")40 [[1, 'a b d'], [10, 'a c d']]41 >>> g("a[b:[c:d:2]:1]e")42 [[1, 'abe'], [2, 'ace'], [10, 'ade']]43 >>> g("a [unbalanced")44 [[10, 'a [unbalanced']]45 >>> g("a [b:.5] c")46 [[5, 'a c'], [10, 'a b c']]47 >>> g("a [{b|d{:.5] c") # not handling this right now48 [[5, 'a c'], [10, 'a {b|d{ c']]49 >>> g("((a][:b:c [d:3]")50 [[3, '((a][:b:c '], [10, '((a][:b:c d']]51 """52 53 def collect_steps(steps, tree):54 l = [steps]55 class CollectSteps(lark.Visitor):56 def scheduled(self, tree):57 tree.children[-1] = float(tree.children[-1])58 if tree.children[-1] < 1:59 tree.children[-1] *= steps60 tree.children[-1] = min(steps, int(tree.children[-1]))61 l.append(tree.children[-1])62 CollectSteps().visit(tree)63 return sorted(set(l))64 65 def at_step(step, tree):66 class AtStep(lark.Transformer):67 def scheduled(self, args):68 before, after, _, when = args69 yield before or () if step <= when else after70 def start(self, args):71 def flatten(x):72 if type(x) == str:73 yield x74 else:75 for gen in x:76 yield from flatten(gen)77 return ''.join(flatten(args))78 def plain(self, args):79 yield args[0].value80 def __default__(self, data, children, meta):81 for child in children:82 yield from child83 return AtStep().transform(tree)84 85 def get_schedule(prompt):86 try:87 tree = schedule_parser.parse(prompt)88 except lark.exceptions.LarkError as e:89 if 0:90 import traceback91 traceback.print_exc()92 return [[steps, prompt]]93 return [[t, at_step(t, tree)] for t in collect_steps(steps, tree)]94 95 promptdict = {prompt: get_schedule(prompt) for prompt in set(prompts)}96 return [promptdict[prompt] for prompt in prompts]97 98 99ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])100 101 102def get_learned_conditioning(model, prompts, steps):103 """converts a list of prompts into a list of prompt schedules - each schedule is a list of ScheduledPromptConditioning, specifying the comdition (cond),104 and the sampling step at which this condition is to be replaced by the next one.105 106 Input:107 (model, ['a red crown', 'a [blue:green:5] jeweled crown'], 20)108 109 Output:110 [111 [112 ScheduledPromptConditioning(end_at_step=20, cond=tensor([[-0.3886, 0.0229, -0.0523, ..., -0.4901, -0.3066, 0.0674], ..., [ 0.3317, -0.5102, -0.4066, ..., 0.4119, -0.7647, -1.0160]], device='cuda:0'))113 ],114 [115 ScheduledPromptConditioning(end_at_step=5, cond=tensor([[-0.3886, 0.0229, -0.0522, ..., -0.4901, -0.3067, 0.0673], ..., [-0.0192, 0.3867, -0.4644, ..., 0.1135, -0.3696, -0.4625]], device='cuda:0')),116 ScheduledPromptConditioning(end_at_step=20, cond=tensor([[-0.3886, 0.0229, -0.0522, ..., -0.4901, -0.3067, 0.0673], ..., [-0.7352, -0.4356, -0.7888, ..., 0.6994, -0.4312, -1.2593]], device='cuda:0'))117 ]118 ]119 """120 res = []121 122 prompt_schedules = get_learned_conditioning_prompt_schedules(prompts, steps)123 cache = {}124 125 for prompt, prompt_schedule in zip(prompts, prompt_schedules):126 127 cached = cache.get(prompt, None)128 if cached is not None:129 res.append(cached)130 continue131 132 texts = [x[1] for x in prompt_schedule]133 conds = model.get_learned_conditioning(texts)134 135 cond_schedule = []136 for i, (end_at_step, text) in enumerate(prompt_schedule):137 cond_schedule.append(ScheduledPromptConditioning(end_at_step, conds[i]))138 139 cache[prompt] = cond_schedule140 res.append(cond_schedule)141 142 return res143 144 145re_AND = re.compile(r"\bAND\b")146re_weight = re.compile(r"^(.*?)(?:\s*:\s*([-+]?(?:\d+\.?|\d*\.\d+)))?\s*$")147 148def get_multicond_prompt_list(prompts):149 res_indexes = []150 151 prompt_flat_list = []152 prompt_indexes = {}153 154 for prompt in prompts:155 subprompts = re_AND.split(prompt)156 157 indexes = []158 for subprompt in subprompts:159 match = re_weight.search(subprompt)160 161 text, weight = match.groups() if match is not None else (subprompt, 1.0)162 163 weight = float(weight) if weight is not None else 1.0164 165 index = prompt_indexes.get(text, None)166 if index is None:167 index = len(prompt_flat_list)168 prompt_flat_list.append(text)169 prompt_indexes[text] = index170 171 indexes.append((index, weight))172 173 res_indexes.append(indexes)174 175 return res_indexes, prompt_flat_list, prompt_indexes176 177 178class ComposableScheduledPromptConditioning:179 def __init__(self, schedules, weight=1.0):180 self.schedules: List[ScheduledPromptConditioning] = schedules181 self.weight: float = weight182 183 184class MulticondLearnedConditioning:185 def __init__(self, shape, batch):186 self.shape: tuple = shape # the shape field is needed to send this object to DDIM/PLMS187 self.batch: List[List[ComposableScheduledPromptConditioning]] = batch188 189def get_multicond_learned_conditioning(model, prompts, steps) -> MulticondLearnedConditioning:190 """same as get_learned_conditioning, but returns a list of ScheduledPromptConditioning along with the weight objects for each prompt.191 For each prompt, the list is obtained by splitting the prompt using the AND separator.192 193 https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/194 """195 196 res_indexes, prompt_flat_list, prompt_indexes = get_multicond_prompt_list(prompts)197 198 learned_conditioning = get_learned_conditioning(model, prompt_flat_list, steps)199 200 res = []201 for indexes in res_indexes:202 res.append([ComposableScheduledPromptConditioning(learned_conditioning[i], weight) for i, weight in indexes])203 204 return MulticondLearnedConditioning(shape=(len(prompts),), batch=res)205 206 207def reconstruct_cond_batch(c: List[List[ScheduledPromptConditioning]], current_step):208 param = c[0][0].cond209 res = torch.zeros((len(c),) + param.shape, device=param.device, dtype=param.dtype)210 for i, cond_schedule in enumerate(c):211 target_index = 0212 for current, (end_at, cond) in enumerate(cond_schedule):213 if current_step <= end_at:214 target_index = current215 break216 res[i] = cond_schedule[target_index].cond217 218 return res219 220 221def reconstruct_multicond_batch(c: MulticondLearnedConditioning, current_step):222 param = c.batch[0][0].schedules[0].cond223 224 tensors = []225 conds_list = []226 227 for batch_no, composable_prompts in enumerate(c.batch):228 conds_for_batch = []229 230 for cond_index, composable_prompt in enumerate(composable_prompts):231 target_index = 0232 for current, (end_at, cond) in enumerate(composable_prompt.schedules):233 if current_step <= end_at:234 target_index = current235 break236 237 conds_for_batch.append((len(tensors), composable_prompt.weight))238 tensors.append(composable_prompt.schedules[target_index].cond)239 240 conds_list.append(conds_for_batch)241 242 return conds_list, torch.stack(tensors).to(device=param.device, dtype=param.dtype)243 244 245re_attention = re.compile(r"""246\\\(|247\\\)|248\\\[|249\\]|250\\\\|251\\|252\(|253\[|254:([+-]?[.\d]+)\)|255\)|256]|257[^\\()\[\]:]+|258:259""", re.X)260 261 262def parse_prompt_attention(text):263 """264 Parses a string with attention tokens and returns a list of pairs: text and its assoicated weight.265 Accepted tokens are:266 (abc) - increases attention to abc by a multiplier of 1.1267 (abc:3.12) - increases attention to abc by a multiplier of 3.12268 [abc] - decreases attention to abc by a multiplier of 1.1269 \( - literal character '('270 \[ - literal character '['271 \) - literal character ')'272 \] - literal character ']'273 \\ - literal character '\'274 anything else - just text275 276 >>> parse_prompt_attention('normal text')277 [['normal text', 1.0]]278 >>> parse_prompt_attention('an (important) word')279 [['an ', 1.0], ['important', 1.1], [' word', 1.0]]280 >>> parse_prompt_attention('(unbalanced')281 [['unbalanced', 1.1]]282 >>> parse_prompt_attention('\(literal\]')283 [['(literal]', 1.0]]284 >>> parse_prompt_attention('(unnecessary)(parens)')285 [['unnecessaryparens', 1.1]]286 >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')287 [['a ', 1.0],288 ['house', 1.5730000000000004],289 [' ', 1.1],290 ['on', 1.0],291 [' a ', 1.1],292 ['hill', 0.55],293 [', sun, ', 1.1],294 ['sky', 1.4641000000000006],295 ['.', 1.1]]296 """297 298 res = []299 round_brackets = []300 square_brackets = []301 302 round_bracket_multiplier = 1.1303 square_bracket_multiplier = 1 / 1.1304 305 def multiply_range(start_position, multiplier):306 for p in range(start_position, len(res)):307 res[p][1] *= multiplier308 309 for m in re_attention.finditer(text):310 text = m.group(0)311 weight = m.group(1)312 313 if text.startswith('\\'):314 res.append([text[1:], 1.0])315 elif text == '(':316 round_brackets.append(len(res))317 elif text == '[':318 square_brackets.append(len(res))319 elif weight is not None and len(round_brackets) > 0:320 multiply_range(round_brackets.pop(), float(weight))321 elif text == ')' and len(round_brackets) > 0:322 multiply_range(round_brackets.pop(), round_bracket_multiplier)323 elif text == ']' and len(square_brackets) > 0:324 multiply_range(square_brackets.pop(), square_bracket_multiplier)325 else:326 res.append([text, 1.0])327 328 for pos in round_brackets:329 multiply_range(pos, round_bracket_multiplier)330 331 for pos in square_brackets:332 multiply_range(pos, square_bracket_multiplier)333 334 if len(res) == 0:335 res = [["", 1.0]]336 337 # merge runs of identical weights338 i = 0339 while i + 1 < len(res):340 if res[i][1] == res[i + 1][1]:341 res[i][0] += res[i + 1][0]342 res.pop(i + 1)343 else:344 i += 1345 346 return res347 348if __name__ == "__main__":349 import doctest350 doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)351else:352 import torch # doctest faster353