cymic/Waifu_Diffusion_Webui

1import math2import os3import sys4import traceback5import torch6import numpy as np7from torch import einsum8from torch.nn.functional import silu9 10import modules.textual_inversion.textual_inversion11from modules import prompt_parser, devices, sd_hijack_optimizations, shared, hypernetwork12from modules.shared import opts, device, cmd_opts13 14import ldm.modules.attention15import ldm.modules.diffusionmodules.model16 17attention_CrossAttention_forward = ldm.modules.attention.CrossAttention.forward18diffusionmodules_model_nonlinearity = ldm.modules.diffusionmodules.model.nonlinearity19diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward20 21 22def apply_optimizations():23    undo_optimizations()24 25    ldm.modules.diffusionmodules.model.nonlinearity = silu26 27    if cmd_opts.opt_split_attention_v1:28        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v129    elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):30        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward31        ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.cross_attention_attnblock_forward32 33 34def undo_optimizations():35    ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward36    ldm.modules.diffusionmodules.model.nonlinearity = diffusionmodules_model_nonlinearity37    ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward38 39 40class StableDiffusionModelHijack:41    fixes = None42    comments = []43    layers = None44    circular_enabled = False45    clip = None46 47    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase(cmd_opts.embeddings_dir)48 49    def hijack(self, m):50        model_embeddings = m.cond_stage_model.transformer.text_model.embeddings51 52        model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)53        m.cond_stage_model = FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)54 55        self.clip = m.cond_stage_model56 57        apply_optimizations()58 59        def flatten(el):60            flattened = [flatten(children) for children in el.children()]61            res = [el]62            for c in flattened:63                res += c64            return res65 66        self.layers = flatten(m)67 68    def undo_hijack(self, m):69        if type(m.cond_stage_model) == FrozenCLIPEmbedderWithCustomWords:70            m.cond_stage_model = m.cond_stage_model.wrapped71 72        model_embeddings = m.cond_stage_model.transformer.text_model.embeddings73        if type(model_embeddings.token_embedding) == EmbeddingsWithFixes:74            model_embeddings.token_embedding = model_embeddings.token_embedding.wrapped75 76    def apply_circular(self, enable):77        if self.circular_enabled == enable:78            return79 80        self.circular_enabled = enable81 82        for layer in [layer for layer in self.layers if type(layer) == torch.nn.Conv2d]:83            layer.padding_mode = 'circular' if enable else 'zeros'84 85    def tokenize(self, text):86        max_length = self.clip.max_length - 287        _, remade_batch_tokens, _, _, _, token_count = self.clip.process_text([text])88        return remade_batch_tokens[0], token_count, max_length89 90 91class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):92    def __init__(self, wrapped, hijack):93        super().__init__()94        self.wrapped = wrapped95        self.hijack: StableDiffusionModelHijack = hijack96        self.tokenizer = wrapped.tokenizer97        self.max_length = wrapped.max_length98        self.token_mults = {}99 100        tokens_with_parens = [(k, v) for k, v in self.tokenizer.get_vocab().items() if '(' in k or ')' in k or '[' in k or ']' in k]101        for text, ident in tokens_with_parens:102            mult = 1.0103            for c in text:104                if c == '[':105                    mult /= 1.1106                if c == ']':107                    mult *= 1.1108                if c == '(':109                    mult *= 1.1110                if c == ')':111                    mult /= 1.1112 113            if mult != 1.0:114                self.token_mults[ident] = mult115 116    def tokenize_line(self, line, used_custom_terms, hijack_comments):117        id_start = self.wrapped.tokenizer.bos_token_id118        id_end = self.wrapped.tokenizer.eos_token_id119        maxlen = self.wrapped.max_length120 121        if opts.enable_emphasis:122            parsed = prompt_parser.parse_prompt_attention(line)123        else:124            parsed = [[line, 1.0]]125 126        tokenized = self.wrapped.tokenizer([text for text, _ in parsed], truncation=False, add_special_tokens=False)["input_ids"]127 128        fixes = []129        remade_tokens = []130        multipliers = []131 132        for tokens, (text, weight) in zip(tokenized, parsed):133            i = 0134            while i < len(tokens):135                token = tokens[i]136 137                embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)138 139                if embedding is None:140                    remade_tokens.append(token)141                    multipliers.append(weight)142                    i += 1143                else:144                    emb_len = int(embedding.vec.shape[0])145                    fixes.append((len(remade_tokens), embedding))146                    remade_tokens += [0] * emb_len147                    multipliers += [weight] * emb_len148                    used_custom_terms.append((embedding.name, embedding.checksum()))149                    i += embedding_length_in_tokens150 151        if len(remade_tokens) > maxlen - 2:152            vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}153            ovf = remade_tokens[maxlen - 2:]154            overflowing_words = [vocab.get(int(x), "") for x in ovf]155            overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))156            hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")157 158        token_count = len(remade_tokens)159        remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))160        remade_tokens = [id_start] + remade_tokens[0:maxlen - 2] + [id_end]161 162        multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))163        multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]164 165        return remade_tokens, fixes, multipliers, token_count166 167    def process_text(self, texts):168        used_custom_terms = []169        remade_batch_tokens = []170        hijack_comments = []171        hijack_fixes = []172        token_count = 0173 174        cache = {}175        batch_multipliers = []176        for line in texts:177            if line in cache:178                remade_tokens, fixes, multipliers = cache[line]179            else:180                remade_tokens, fixes, multipliers, token_count = self.tokenize_line(line, used_custom_terms, hijack_comments)181 182                cache[line] = (remade_tokens, fixes, multipliers)183 184            remade_batch_tokens.append(remade_tokens)185            hijack_fixes.append(fixes)186            batch_multipliers.append(multipliers)187 188        return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count189 190 191    def process_text_old(self, text):192        id_start = self.wrapped.tokenizer.bos_token_id193        id_end = self.wrapped.tokenizer.eos_token_id194        maxlen = self.wrapped.max_length195        used_custom_terms = []196        remade_batch_tokens = []197        overflowing_words = []198        hijack_comments = []199        hijack_fixes = []200        token_count = 0201 202        cache = {}203        batch_tokens = self.wrapped.tokenizer(text, truncation=False, add_special_tokens=False)["input_ids"]204        batch_multipliers = []205        for tokens in batch_tokens:206            tuple_tokens = tuple(tokens)207 208            if tuple_tokens in cache:209                remade_tokens, fixes, multipliers = cache[tuple_tokens]210            else:211                fixes = []212                remade_tokens = []213                multipliers = []214                mult = 1.0215 216                i = 0217                while i < len(tokens):218                    token = tokens[i]219 220                    embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)221 222                    mult_change = self.token_mults.get(token) if opts.enable_emphasis else None223                    if mult_change is not None:224                        mult *= mult_change225                        i += 1226                    elif embedding is None:227                        remade_tokens.append(token)228                        multipliers.append(mult)229                        i += 1230                    else:231                        emb_len = int(embedding.vec.shape[0])232                        fixes.append((len(remade_tokens), embedding))233                        remade_tokens += [0] * emb_len234                        multipliers += [mult] * emb_len235                        used_custom_terms.append((embedding.name, embedding.checksum()))236                        i += embedding_length_in_tokens237 238                if len(remade_tokens) > maxlen - 2:239                    vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}240                    ovf = remade_tokens[maxlen - 2:]241                    overflowing_words = [vocab.get(int(x), "") for x in ovf]242                    overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))243                    hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")244 245                token_count = len(remade_tokens)246                remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))247                remade_tokens = [id_start] + remade_tokens[0:maxlen-2] + [id_end]248                cache[tuple_tokens] = (remade_tokens, fixes, multipliers)249 250            multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))251            multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]252 253            remade_batch_tokens.append(remade_tokens)254            hijack_fixes.append(fixes)255            batch_multipliers.append(multipliers)256        return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count257 258    def forward(self, text):259 260        if opts.use_old_emphasis_implementation:261            batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text_old(text)262        else:263            batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = self.process_text(text)264 265        self.hijack.fixes = hijack_fixes266        self.hijack.comments = hijack_comments267 268        if len(used_custom_terms) > 0:269            self.hijack.comments.append("Used embeddings: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms]))270 271        tokens = torch.asarray(remade_batch_tokens).to(device)272        outputs = self.wrapped.transformer(input_ids=tokens)273        z = outputs.last_hidden_state274 275        # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise276        batch_multipliers = torch.asarray(batch_multipliers).to(device)277        original_mean = z.mean()278        z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)279        new_mean = z.mean()280        z *= original_mean / new_mean281 282        return z283 284 285class EmbeddingsWithFixes(torch.nn.Module):286    def __init__(self, wrapped, embeddings):287        super().__init__()288        self.wrapped = wrapped289        self.embeddings = embeddings290 291    def forward(self, input_ids):292        batch_fixes = self.embeddings.fixes293        self.embeddings.fixes = None294 295        inputs_embeds = self.wrapped(input_ids)296 297        if batch_fixes is None or len(batch_fixes) == 0 or max([len(x) for x in batch_fixes]) == 0:298            return inputs_embeds299 300        vecs = []301        for fixes, tensor in zip(batch_fixes, inputs_embeds):302            for offset, embedding in fixes:303                emb = embedding.vec304                emb_len = min(tensor.shape[0]-offset-1, emb.shape[0])305                tensor = torch.cat([tensor[0:offset+1], emb[0:emb_len], tensor[offset+1+emb_len:]])306 307            vecs.append(tensor)308 309        return torch.stack(vecs)310 311 312def add_circular_option_to_conv_2d():313    conv2d_constructor = torch.nn.Conv2d.__init__314 315    def conv2d_constructor_circular(self, *args, **kwargs):316        return conv2d_constructor(self, *args, padding_mode='circular', **kwargs)317 318    torch.nn.Conv2d.__init__ = conv2d_constructor_circular319 320 321model_hijack = StableDiffusionModelHijack()322