ethansmith2000/MegaEdit
0
1import re2import torch3 4re_attention = re.compile(5 r"""6\\\(|7\\\)|8\\\[|9\\]|10\\\\|11\\|12\(|13\[|14:([+-]?[.\d]+)\)|15\)|16]|17[^\\()\[\]:]+|18:19""",20 re.X,21)22 23 24def parse_prompt_attention(text):25 """26 Parses a string with attention tokens and returns a list of pairs: text and its associated weight.27 Accepted tokens are:28 (abc) - increases attention to abc by a multiplier of 1.129 (abc:3.12) - increases attention to abc by a multiplier of 3.1230 [abc] - decreases attention to abc by a multiplier of 1.131 \( - literal character '('32 \[ - literal character '['33 \) - literal character ')'34 \] - literal character ']'35 \\ - literal character '\'36 anything else - just text37 >>> parse_prompt_attention('normal text')38 [['normal text', 1.0]]39 >>> parse_prompt_attention('an (important) word')40 [['an ', 1.0], ['important', 1.1], [' word', 1.0]]41 >>> parse_prompt_attention('(unbalanced')42 [['unbalanced', 1.1]]43 >>> parse_prompt_attention('\(literal\]')44 [['(literal]', 1.0]]45 >>> parse_prompt_attention('(unnecessary)(parens)')46 [['unnecessaryparens', 1.1]]47 >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')48 [['a ', 1.0],49 ['house', 1.5730000000000004],50 [' ', 1.1],51 ['on', 1.0],52 [' a ', 1.1],53 ['hill', 0.55],54 [', sun, ', 1.1],55 ['sky', 1.4641000000000006],56 ['.', 1.1]]57 """58 59 res = []60 round_brackets = []61 square_brackets = []62 63 round_bracket_multiplier = 1.164 square_bracket_multiplier = 1 / 1.165 66 def multiply_range(start_position, multiplier):67 for p in range(start_position, len(res)):68 res[p][1] *= multiplier69 70 for m in re_attention.finditer(text):71 text = m.group(0)72 weight = m.group(1)73 74 if text.startswith("\\"):75 res.append([text[1:], 1.0])76 elif text == "(":77 round_brackets.append(len(res))78 elif text == "[":79 square_brackets.append(len(res))80 elif weight is not None and len(round_brackets) > 0:81 multiply_range(round_brackets.pop(), float(weight))82 elif text == ")" and len(round_brackets) > 0:83 multiply_range(round_brackets.pop(), round_bracket_multiplier)84 elif text == "]" and len(square_brackets) > 0:85 multiply_range(square_brackets.pop(), square_bracket_multiplier)86 else:87 res.append([text, 1.0])88 89 for pos in round_brackets:90 multiply_range(pos, round_bracket_multiplier)91 92 for pos in square_brackets:93 multiply_range(pos, square_bracket_multiplier)94 95 if len(res) == 0:96 res = [["", 1.0]]97 98 # merge runs of identical weights99 i = 0100 while i + 1 < len(res):101 if res[i][1] == res[i + 1][1]:102 res[i][0] += res[i + 1][0]103 res.pop(i + 1)104 else:105 i += 1106 107 return res108 109 110def get_prompts_with_weights(pipe, prompt, max_length):111 r"""112 Tokenize a list of prompts and return its tokens with weights of each token.113 No padding, starting or ending token is included.114 """115 assert isinstance(prompt, list)116 tokens = []117 weights = []118 truncated = False119 for text in prompt:120 texts_and_weights = parse_prompt_attention(text)121 text_token = []122 text_weight = []123 for word, weight in texts_and_weights:124 # tokenize and discard the starting and the ending token125 token = pipe.tokenizer(word).input_ids[1:-1]126 text_token += token127 # copy the weight by length of token128 text_weight += [weight] * len(token)129 # stop if the text is too long (longer than truncation limit)130 if len(text_token) > max_length:131 truncated = True132 break133 # truncate134 if len(text_token) > max_length:135 truncated = True136 text_token = text_token[:max_length]137 text_weight = text_weight[:max_length]138 tokens.append(text_token)139 weights.append(text_weight)140 141 return tokens, weights142 143 144def pad_tokens_and_weights(tokens, weights, max_length, bos, eos, no_boseos_middle=True, chunk_length=77):145 r"""146 Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length.147 """148 max_embeddings_multiples = (max_length - 2) // (chunk_length - 2)149 weights_length = max_length if no_boseos_middle else max_embeddings_multiples * chunk_length150 for i in range(len(tokens)):151 tokens[i] = [bos] + tokens[i] + [eos] * (max_length - 1 - len(tokens[i]))152 if no_boseos_middle:153 weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i]))154 else:155 w = []156 if len(weights[i]) == 0:157 w = [1.0] * weights_length158 else:159 for j in range(max_embeddings_multiples):160 w.append(1.0) # weight for starting token in this chunk161 w += weights[i][j * (chunk_length - 2): min(len(weights[i]), (j + 1) * (chunk_length - 2))]162 w.append(1.0) # weight for ending token in this chunk163 w += [1.0] * (weights_length - len(w))164 weights[i] = w[:]165 166 return tokens, weights167 168 169def get_unweighted_text_embeddings(170 pipe,171 text_input,172 chunk_length,173 no_boseos_middle=True,174):175 """176 When the length of tokens is a multiple of the capacity of the text encoder,177 it should be split into chunks and sent to the text encoder individually.178 """179 max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2)180 if max_embeddings_multiples > 1:181 text_embeddings = []182 for i in range(max_embeddings_multiples):183 # extract the i-th chunk184 text_input_chunk = text_input[:, i * (chunk_length - 2): (i + 1) * (chunk_length - 2) + 2].clone()185 186 # cover the head and the tail by the starting and the ending tokens187 text_input_chunk[:, 0] = text_input[0, 0]188 text_input_chunk[:, -1] = text_input[0, -1]189 text_embedding = pipe.text_encoder(text_input_chunk).last_hidden_state190 191 if no_boseos_middle:192 if i == 0:193 # discard the ending token194 text_embedding = text_embedding[:, :-1]195 elif i == max_embeddings_multiples - 1:196 # discard the starting token197 text_embedding = text_embedding[:, 1:]198 else:199 # discard both starting and ending tokens200 text_embedding = text_embedding[:, 1:-1]201 202 text_embeddings.append(text_embedding)203 text_embeddings = torch.concat(text_embeddings, axis=1)204 else:205 text_embeddings = pipe.text_encoder(text_input).last_hidden_state206 return text_embeddings207 208 209def get_weighted_text_embeddings(210 pipe,211 prompt,212 max_embeddings_multiples=1,213 no_boseos_middle=False,214 skip_parsing=False,215 skip_weighting=False,216 **kwargs,217):218 r"""219 Prompts can be assigned with local weights using brackets. For example,220 prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful',221 and the embedding tokens corresponding to the words get multiplied by a constant, 1.1.222 Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean.223 Args:224 pipe (`StableDiffusionPipeline`):225 Pipe to provide access to the tokenizer and the text encoder.226 prompt (`str` or `List[str]`):227 The prompt or prompts to guide the image generation.228 uncond_prompt (`str` or `List[str]`):229 The unconditional prompt or prompts for guide the image generation. If unconditional prompt230 is provided, the embeddings of prompt and uncond_prompt are concatenated.231 max_embeddings_multiples (`int`, *optional*, defaults to `3`):232 The max multiple length of prompt embeddings compared to the max output length of text encoder.233 no_boseos_middle (`bool`, *optional*, defaults to `False`):234 If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and235 ending token in each of the chunk in the middle.236 skip_parsing (`bool`, *optional*, defaults to `False`):237 Skip the parsing of brackets.238 skip_weighting (`bool`, *optional*, defaults to `False`):239 Skip the weighting. When the parsing is skipped, it is forced True.240 """241 max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2242 if isinstance(prompt, str):243 prompt = [prompt]244 245 if not skip_parsing:246 prompt_tokens, prompt_weights = get_prompts_with_weights(pipe, prompt, max_length - 2)247 else:248 prompt_tokens = [249 token[1:-1] for token in pipe.tokenizer(prompt, max_length=max_length, truncation=True).input_ids250 ]251 prompt_weights = [[1.0] * len(token) for token in prompt_tokens]252 253 # round up the longest length of tokens to a multiple of (model_max_length - 2)254 max_length = max([len(token) for token in prompt_tokens])255 256 max_embeddings_multiples = min(257 max_embeddings_multiples,258 (max_length - 1) // (pipe.tokenizer.model_max_length - 2) + 1,259 )260 max_embeddings_multiples = max(1, max_embeddings_multiples)261 max_length = (pipe.tokenizer.model_max_length - 2) * max_embeddings_multiples + 2262 263 # pad the length of tokens and weights264 bos = pipe.tokenizer.bos_token_id265 eos = pipe.tokenizer.eos_token_id266 prompt_tokens, prompt_weights = pad_tokens_and_weights(267 prompt_tokens,268 prompt_weights,269 max_length,270 bos,271 eos,272 no_boseos_middle=no_boseos_middle,273 chunk_length=pipe.tokenizer.model_max_length,274 )275 prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=pipe.device)276 277 # get the embeddings278 text_embeddings = get_unweighted_text_embeddings(279 pipe,280 prompt_tokens,281 pipe.tokenizer.model_max_length,282 no_boseos_middle=no_boseos_middle,283 )284 prompt_weights = torch.tensor(prompt_weights, dtype=text_embeddings.dtype, device=pipe.device)285 286 # assign weights to the prompts and normalize in the sense of mean287 # TODO: should we normalize by chunk or in a whole (current implementation)?288 if (not skip_parsing) and (not skip_weighting):289 previous_mean = text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype)290 text_embeddings *= prompt_weights.unsqueeze(-1)291 current_mean = text_embeddings.float().mean(axis=[-2, -1]).to(text_embeddings.dtype)292 text_embeddings *= (previous_mean / current_mean).unsqueeze(-1).unsqueeze(-1)293 294 return text_embeddings, prompt_tokens, prompt_weights295