sd1_clip.py 20.5 KB
Newer Older
comfyanonymous's avatar
comfyanonymous committed
1
2
import os

3
from transformers import CLIPTokenizer
4
import comfy.ops
comfyanonymous's avatar
comfyanonymous committed
5
import torch
6
import traceback
7
import zipfile
8
from . import model_management
9
10
import comfy.clip_model
import json
11
import logging
comfyanonymous's avatar
comfyanonymous committed
12

13
14
15
16
17
18
19
20
21
22
23
24
def gen_empty_tokens(special_tokens, length):
    start_token = special_tokens.get("start", None)
    end_token = special_tokens.get("end", None)
    pad_token = special_tokens.get("pad")
    output = []
    if start_token is not None:
        output.append(start_token)
    if end_token is not None:
        output.append(end_token)
    output += [pad_token] * (length - len(output))
    return output

comfyanonymous's avatar
comfyanonymous committed
25
26
class ClipTokenWeightEncoder:
    def encode_token_weights(self, token_weight_pairs):
27
28
29
        to_encode = list()
        max_token_len = 0
        has_weights = False
comfyanonymous's avatar
comfyanonymous committed
30
        for x in token_weight_pairs:
31
            tokens = list(map(lambda a: a[0], x))
32
33
            max_token_len = max(len(tokens), max_token_len)
            has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
34
35
            to_encode.append(tokens)

36
37
38
39
        sections = len(to_encode)
        if has_weights or sections == 0:
            to_encode.append(gen_empty_tokens(self.special_tokens, max_token_len))

40
        out, pooled = self.encode(to_encode)
41
        if pooled is not None:
42
            first_pooled = pooled[0:1].to(model_management.intermediate_device())
43
        else:
44
            first_pooled = pooled
45
46

        output = []
47
        for k in range(0, sections):
48
            z = out[k:k+1]
49
50
51
52
53
54
55
            if has_weights:
                z_empty = out[-1]
                for i in range(len(z)):
                    for j in range(len(z[i])):
                        weight = token_weight_pairs[k][j][1]
                        if weight != 1.0:
                            z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
56
57
            output.append(z)

comfyanonymous's avatar
comfyanonymous committed
58
        if (len(output) == 0):
59
60
            return out[-1:].to(model_management.intermediate_device()), first_pooled
        return torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled
comfyanonymous's avatar
comfyanonymous committed
61

62
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
comfyanonymous's avatar
comfyanonymous committed
63
64
65
66
67
68
69
    """Uses the CLIP transformer encoder for text (from huggingface)"""
    LAYERS = [
        "last",
        "pooled",
        "hidden"
    ]
    def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
70
                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=comfy.clip_model.CLIPTextModel,
71
72
                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, zero_out_masked=False,
                 return_projected_pooled=True):  # clip-vit-base-patch32
comfyanonymous's avatar
comfyanonymous committed
73
74
        super().__init__()
        assert layer in self.LAYERS
75
76
77
78
79
80
81

        if textmodel_json_config is None:
            textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")

        with open(textmodel_json_config) as f:
            config = json.load(f)

82
        self.transformer = model_class(config, dtype, device, comfy.ops.manual_cast)
83
        self.num_layers = self.transformer.num_layers
84

comfyanonymous's avatar
comfyanonymous committed
85
86
87
88
89
        self.max_length = max_length
        if freeze:
            self.freeze()
        self.layer = layer
        self.layer_idx = None
90
        self.special_tokens = special_tokens
91

92
        self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
93
        self.enable_attention_masks = enable_attention_masks
94
        self.zero_out_masked = zero_out_masked
95

96
        self.layer_norm_hidden_state = layer_norm_hidden_state
97
        self.return_projected_pooled = return_projected_pooled
98

comfyanonymous's avatar
comfyanonymous committed
99
100
        if layer == "hidden":
            assert layer_idx is not None
101
            assert abs(layer_idx) < self.num_layers
102
103
            self.set_clip_options({"layer": layer_idx})
        self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
comfyanonymous's avatar
comfyanonymous committed
104
105
106
107
108
109
110

    def freeze(self):
        self.transformer = self.transformer.eval()
        #self.train = disabled_train
        for param in self.parameters():
            param.requires_grad = False

111
112
113
114
    def set_clip_options(self, options):
        layer_idx = options.get("layer", self.layer_idx)
        self.return_projected_pooled = options.get("projected_pooled", self.return_projected_pooled)
        if layer_idx is None or abs(layer_idx) > self.num_layers:
comfyanonymous's avatar
comfyanonymous committed
115
116
117
118
119
            self.layer = "last"
        else:
            self.layer = "hidden"
            self.layer_idx = layer_idx

120
121
122
123
    def reset_clip_options(self):
        self.layer = self.options_default[0]
        self.layer_idx = self.options_default[1]
        self.return_projected_pooled = self.options_default[2]
124

125
126
    def set_up_textual_embeddings(self, tokens, current_embeds):
        out_tokens = []
127
        next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
128
129
130
131
132
133
        embedding_weights = []

        for x in tokens:
            tokens_temp = []
            for y in x:
                if isinstance(y, int):
134
135
                    if y == token_dict_size: #EOS token
                        y = -1
136
137
                    tokens_temp += [y]
                else:
138
139
140
141
142
                    if y.shape[0] == current_embeds.weight.shape[1]:
                        embedding_weights += [y]
                        tokens_temp += [next_new_token]
                        next_new_token += 1
                    else:
143
                        logging.warning("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored {} != {}".format(y.shape[0], current_embeds.weight.shape[1]))
144
            while len(tokens_temp) < len(x):
145
                tokens_temp += [self.special_tokens["pad"]]
146
147
            out_tokens += [tokens_temp]

148
        n = token_dict_size
149
        if len(embedding_weights) > 0:
150
151
            new_embedding = torch.nn.Embedding(next_new_token + 1, current_embeds.weight.shape[1], device=current_embeds.weight.device, dtype=current_embeds.weight.dtype)
            new_embedding.weight[:token_dict_size] = current_embeds.weight[:-1]
152
153
154
            for x in embedding_weights:
                new_embedding.weight[n] = x
                n += 1
155
            new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
156
            self.transformer.set_input_embeddings(new_embedding)
157
158
159
160
161
162

        processed_tokens = []
        for x in out_tokens:
            processed_tokens += [list(map(lambda a: n if a == -1 else a, x))] #The EOS token should always be the largest one

        return processed_tokens
163

comfyanonymous's avatar
comfyanonymous committed
164
    def forward(self, tokens):
165
        backup_embeds = self.transformer.get_input_embeddings()
166
        device = backup_embeds.weight.device
167
        tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
168
169
        tokens = torch.LongTensor(tokens).to(device)

170
171
172
        attention_mask = None
        if self.enable_attention_masks:
            attention_mask = torch.zeros_like(tokens)
173
            end_token = self.special_tokens.get("end", -1)
174
175
176
            for x in range(attention_mask.shape[0]):
                for y in range(attention_mask.shape[1]):
                    attention_mask[x, y] = 1
177
                    if tokens[x, y] == end_token:
178
179
180
181
182
183
                        break

        outputs = self.transformer(tokens, attention_mask, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
        self.transformer.set_input_embeddings(backup_embeds)

        if self.layer == "last":
184
            z = outputs[0].float()
comfyanonymous's avatar
comfyanonymous committed
185
        else:
186
187
188
189
            z = outputs[1].float()

        if self.zero_out_masked and attention_mask is not None:
            z *= attention_mask.unsqueeze(-1).float()
190

191
192
193
194
195
196
        pooled_output = None
        if len(outputs) >= 3:
            if not self.return_projected_pooled and len(outputs) >= 4 and outputs[3] is not None:
                pooled_output = outputs[3].float()
            elif outputs[2] is not None:
                pooled_output = outputs[2].float()
197

198
        return z, pooled_output
comfyanonymous's avatar
comfyanonymous committed
199
200
201
202

    def encode(self, tokens):
        return self(tokens)

203
204
205
    def load_sd(self, sd):
        return self.transformer.load_state_dict(sd, strict=False)

comfyanonymous's avatar
comfyanonymous committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def parse_parentheses(string):
    result = []
    current_item = ""
    nesting_level = 0
    for char in string:
        if char == "(":
            if nesting_level == 0:
                if current_item:
                    result.append(current_item)
                    current_item = "("
                else:
                    current_item = "("
            else:
                current_item += char
            nesting_level += 1
        elif char == ")":
            nesting_level -= 1
            if nesting_level == 0:
                result.append(current_item + ")")
                current_item = ""
            else:
                current_item += char
        else:
            current_item += char
    if current_item:
        result.append(current_item)
    return result

def token_weights(string, current_weight):
    a = parse_parentheses(string)
    out = []
    for x in a:
        weight = current_weight
        if len(x) >= 2 and x[-1] == ')' and x[0] == '(':
            x = x[1:-1]
            xx = x.rfind(":")
            weight *= 1.1
            if xx > 0:
                try:
                    weight = float(x[xx+1:])
                    x = x[:xx]
                except:
                    pass
            out += token_weights(x, weight)
        else:
            out += [(x, current_weight)]
    return out

def escape_important(text):
    text = text.replace("\\)", "\0\1")
    text = text.replace("\\(", "\0\2")
    return text

def unescape_important(text):
    text = text.replace("\0\1", ")")
    text = text.replace("\0\2", "(")
    return text

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def safe_load_embed_zip(embed_path):
    with zipfile.ZipFile(embed_path) as myzip:
        names = list(filter(lambda a: "data/" in a, myzip.namelist()))
        names.reverse()
        for n in names:
            with myzip.open(n) as myfile:
                data = myfile.read()
                number = len(data) // 4
                length_embed = 1024 #sd2.x
                if number < 768:
                    continue
                if number % 768 == 0:
                    length_embed = 768 #sd1.x
                num_embeds = number // length_embed
                embed = torch.frombuffer(data, dtype=torch.float)
                out = embed.reshape((num_embeds, length_embed)).clone()
                del embed
                return out

283
284
285
286
287
288
289
def expand_directory_list(directories):
    dirs = set()
    for x in directories:
        dirs.add(x)
        for root, subdir, file in os.walk(x, followlinks=True):
            dirs.add(root)
    return list(dirs)
290

291
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
292
293
294
    if isinstance(embedding_directory, str):
        embedding_directory = [embedding_directory]

295
296
    embedding_directory = expand_directory_list(embedding_directory)

297
298
    valid_file = None
    for embed_dir in embedding_directory:
299
300
301
302
303
304
305
        embed_path = os.path.abspath(os.path.join(embed_dir, embedding_name))
        embed_dir = os.path.abspath(embed_dir)
        try:
            if os.path.commonpath((embed_dir, embed_path)) != embed_dir:
                continue
        except:
            continue
306
307
308
309
310
311
312
        if not os.path.isfile(embed_path):
            extensions = ['.safetensors', '.pt', '.bin']
            for x in extensions:
                t = embed_path + x
                if os.path.isfile(t):
                    valid_file = t
                    break
313
        else:
314
315
316
317
318
319
320
321
            valid_file = embed_path
        if valid_file is not None:
            break

    if valid_file is None:
        return None

    embed_path = valid_file
322

323
324
    embed_out = None

325
326
327
328
    try:
        if embed_path.lower().endswith(".safetensors"):
            import safetensors.torch
            embed = safetensors.torch.load_file(embed_path, device="cpu")
comfyanonymous's avatar
comfyanonymous committed
329
        else:
330
            if 'weights_only' in torch.load.__code__.co_varnames:
331
332
333
334
                try:
                    embed = torch.load(embed_path, weights_only=True, map_location="cpu")
                except:
                    embed_out = safe_load_embed_zip(embed_path)
335
336
337
            else:
                embed = torch.load(embed_path, map_location="cpu")
    except Exception as e:
338
        logging.warning("{}\n\nerror loading embedding, skipping loading: {}".format(traceback.format_exc(), embedding_name))
339
340
        return None

341
342
343
    if embed_out is None:
        if 'string_to_param' in embed:
            values = embed['string_to_param'].values()
344
345
346
347
348
349
350
351
352
353
            embed_out = next(iter(values))
        elif isinstance(embed, list):
            out_list = []
            for x in range(len(embed)):
                for k in embed[x]:
                    t = embed[x][k]
                    if t.shape[-1] != embedding_size:
                        continue
                    out_list.append(t.reshape(-1, t.shape[-1]))
            embed_out = torch.cat(out_list, dim=0)
354
355
        elif embed_key is not None and embed_key in embed:
            embed_out = embed[embed_key]
356
357
        else:
            values = embed.values()
358
            embed_out = next(iter(values))
359
    return embed_out
360

361
class SDTokenizer:
362
    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True, min_length=None):
comfyanonymous's avatar
comfyanonymous committed
363
364
        if tokenizer_path is None:
            tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
365
        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
comfyanonymous's avatar
comfyanonymous committed
366
        self.max_length = max_length
367
        self.min_length = min_length
368

comfyanonymous's avatar
comfyanonymous committed
369
        empty = self.tokenizer('')["input_ids"]
370
371
372
373
374
375
376
377
        if has_start_token:
            self.tokens_start = 1
            self.start_token = empty[0]
            self.end_token = empty[1]
        else:
            self.tokens_start = 0
            self.start_token = None
            self.end_token = empty[0]
comfyanonymous's avatar
comfyanonymous committed
378
        self.pad_with_end = pad_with_end
379
380
        self.pad_to_max_length = pad_to_max_length

comfyanonymous's avatar
comfyanonymous committed
381
382
        vocab = self.tokenizer.get_vocab()
        self.inv_vocab = {v: k for k, v in vocab.items()}
383
384
        self.embedding_directory = embedding_directory
        self.max_word_length = 8
385
        self.embedding_identifier = "embedding:"
386
        self.embedding_size = embedding_size
387
        self.embedding_key = embedding_key
388

389
    def _try_get_embedding(self, embedding_name:str):
390
391
392
393
        '''
        Takes a potential embedding name and tries to retrieve it.
        Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
        '''
394
        embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
395
396
397
        if embed is None:
            stripped = embedding_name.strip(',')
            if len(stripped) < len(embedding_name):
398
                embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
399
400
401
402
                return (embed, embedding_name[len(stripped):])
        return (embed, "")


403
    def tokenize_with_weights(self, text:str, return_word_ids=False):
404
405
406
407
408
409
        '''
        Takes a prompt and converts it to a list of (token, weight, word id) elements.
        Tokens can both be integer tokens and pre computed CLIP tensors.
        Word id values are unique per word and embedding, where the id 0 is reserved for non word tokens.
        Returned list has the dimensions NxM where M is the input size of CLIP
        '''
BlenderNeko's avatar
BlenderNeko committed
410
411
412
413
        if self.pad_with_end:
            pad_token = self.end_token
        else:
            pad_token = 0
comfyanonymous's avatar
comfyanonymous committed
414
415
416
417

        text = escape_important(text)
        parsed_weights = token_weights(text, 1.0)

418
        #tokenize words
comfyanonymous's avatar
comfyanonymous committed
419
        tokens = []
420
421
422
423
424
425
        for weighted_segment, weight in parsed_weights:
            to_tokenize = unescape_important(weighted_segment).replace("\n", " ").split(' ')
            to_tokenize = [x for x in to_tokenize if x != ""]
            for word in to_tokenize:
                #if we find an embedding, deal with the embedding
                if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
426
427
                    embedding_name = word[len(self.embedding_identifier):].strip('\n')
                    embed, leftover = self._try_get_embedding(embedding_name)
428
                    if embed is None:
429
                        logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
430
                    else:
431
                        if len(embed.shape) == 1:
432
                            tokens.append([(embed, weight)])
433
                        else:
434
435
436
437
                            tokens.append([(embed[x], weight) for x in range(embed.shape[0])])
                    #if we accidentally have leftover text, continue parsing using leftover, else move on to next word
                    if leftover != "":
                        word = leftover
438
                    else:
439
440
                        continue
                #parse word
441
                tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
442

443
444
        #reshape token array to CLIP input size
        batched_tokens = []
445
446
447
        batch = []
        if self.start_token is not None:
            batch.append((self.start_token, 1.0, 0))
448
449
        batched_tokens.append(batch)
        for i, t_group in enumerate(tokens):
450
451
            #determine if we're going to try and keep the tokens in a single batch
            is_large = len(t_group) >= self.max_word_length
BlenderNeko's avatar
BlenderNeko committed
452

453
            while len(t_group) > 0:
BlenderNeko's avatar
BlenderNeko committed
454
455
456
                if len(t_group) + len(batch) > self.max_length - 1:
                    remaining_length = self.max_length - len(batch) - 1
                    #break word in two and add end token
457
458
                    if is_large:
                        batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
BlenderNeko's avatar
BlenderNeko committed
459
                        batch.append((self.end_token, 1.0, 0))
460
                        t_group = t_group[remaining_length:]
BlenderNeko's avatar
BlenderNeko committed
461
                    #add end token and pad
462
                    else:
BlenderNeko's avatar
BlenderNeko committed
463
                        batch.append((self.end_token, 1.0, 0))
464
465
                        if self.pad_to_max_length:
                            batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
BlenderNeko's avatar
BlenderNeko committed
466
                    #start new batch
467
468
469
                    batch = []
                    if self.start_token is not None:
                        batch.append((self.start_token, 1.0, 0))
470
                    batched_tokens.append(batch)
471
                else:
472
473
                    batch.extend([(t,w,i+1) for t,w in t_group])
                    t_group = []
474

475
        #fill last batch
476
477
478
        batch.append((self.end_token, 1.0, 0))
        if self.pad_to_max_length:
            batch.extend([(pad_token, 1.0, 0)] * (self.max_length - len(batch)))
479
480
        if self.min_length is not None and len(batch) < self.min_length:
            batch.extend([(pad_token, 1.0, 0)] * (self.min_length - len(batch)))
comfyanonymous's avatar
comfyanonymous committed
481

482
483
        if not return_word_ids:
            batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
comfyanonymous's avatar
comfyanonymous committed
484

485
        return batched_tokens
comfyanonymous's avatar
comfyanonymous committed
486
487
488
489


    def untokenize(self, token_weight_pair):
        return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507


class SD1Tokenizer:
    def __init__(self, embedding_directory=None, clip_name="l", tokenizer=SDTokenizer):
        self.clip_name = clip_name
        self.clip = "clip_{}".format(self.clip_name)
        setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory))

    def tokenize_with_weights(self, text:str, return_word_ids=False):
        out = {}
        out[self.clip_name] = getattr(self, self.clip).tokenize_with_weights(text, return_word_ids)
        return out

    def untokenize(self, token_weight_pair):
        return getattr(self, self.clip).untokenize(token_weight_pair)


class SD1ClipModel(torch.nn.Module):
508
    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
509
510
511
        super().__init__()
        self.clip_name = clip_name
        self.clip = "clip_{}".format(self.clip_name)
512
        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
513

514
515
516
517
        self.dtypes = set()
        if dtype is not None:
            self.dtypes.add(dtype)

518
519
    def set_clip_options(self, options):
        getattr(self, self.clip).set_clip_options(options)
520

521
522
    def reset_clip_options(self):
        getattr(self, self.clip).reset_clip_options()
523
524
525
526
527
528
529
530

    def encode_token_weights(self, token_weight_pairs):
        token_weight_pairs = token_weight_pairs[self.clip_name]
        out, pooled = getattr(self, self.clip).encode_token_weights(token_weight_pairs)
        return out, pooled

    def load_sd(self, sd):
        return getattr(self, self.clip).load_sd(sd)