sd1_clip.py 21.3 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
Mario Klingemann's avatar
Mario Klingemann committed
12
import numbers
comfyanonymous's avatar
comfyanonymous committed
13

14
15
16
17
18
19
20
21
22
23
24
25
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
26
27
class ClipTokenWeightEncoder:
    def encode_token_weights(self, token_weight_pairs):
28
29
30
        to_encode = list()
        max_token_len = 0
        has_weights = False
comfyanonymous's avatar
comfyanonymous committed
31
        for x in token_weight_pairs:
32
            tokens = list(map(lambda a: a[0], x))
33
34
            max_token_len = max(len(tokens), max_token_len)
            has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
35
36
            to_encode.append(tokens)

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

41
42
43
        o = self.encode(to_encode)
        out, pooled = o[:2]

44
        if pooled is not None:
45
            first_pooled = pooled[0:1].to(model_management.intermediate_device())
46
        else:
47
            first_pooled = pooled
48
49

        output = []
50
        for k in range(0, sections):
51
            z = out[k:k+1]
52
53
54
55
56
57
58
            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]
59
60
            output.append(z)

comfyanonymous's avatar
comfyanonymous committed
61
        if (len(output) == 0):
62
63
64
65
66
            r = (out[-1:].to(model_management.intermediate_device()), first_pooled)
        else:
            r = (torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled)
        r = r + tuple(map(lambda a: a[:sections].flatten().unsqueeze(dim=0).to(model_management.intermediate_device()), o[2:]))
        return r
comfyanonymous's avatar
comfyanonymous committed
67

68
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
comfyanonymous's avatar
comfyanonymous committed
69
70
71
72
73
74
75
    """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,
76
                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=comfy.clip_model.CLIPTextModel,
77
                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, zero_out_masked=False,
78
                 return_projected_pooled=True, return_attention_masks=False):  # clip-vit-base-patch32
comfyanonymous's avatar
comfyanonymous committed
79
80
        super().__init__()
        assert layer in self.LAYERS
81
82
83
84
85
86
87

        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)

88
        self.transformer = model_class(config, dtype, device, comfy.ops.manual_cast)
89
        self.num_layers = self.transformer.num_layers
90

comfyanonymous's avatar
comfyanonymous committed
91
92
93
94
95
        self.max_length = max_length
        if freeze:
            self.freeze()
        self.layer = layer
        self.layer_idx = None
96
        self.special_tokens = special_tokens
97

98
        self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
99
        self.enable_attention_masks = enable_attention_masks
100
        self.zero_out_masked = zero_out_masked
101

102
        self.layer_norm_hidden_state = layer_norm_hidden_state
103
        self.return_projected_pooled = return_projected_pooled
104
        self.return_attention_masks = return_attention_masks
105

comfyanonymous's avatar
comfyanonymous committed
106
107
        if layer == "hidden":
            assert layer_idx is not None
108
            assert abs(layer_idx) < self.num_layers
109
110
            self.set_clip_options({"layer": layer_idx})
        self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
comfyanonymous's avatar
comfyanonymous committed
111
112
113
114
115
116
117

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

118
119
120
121
    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
122
123
124
125
126
            self.layer = "last"
        else:
            self.layer = "hidden"
            self.layer_idx = layer_idx

127
128
129
130
    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]
131

132
133
    def set_up_textual_embeddings(self, tokens, current_embeds):
        out_tokens = []
134
        next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
135
136
137
138
139
        embedding_weights = []

        for x in tokens:
            tokens_temp = []
            for y in x:
Mario Klingemann's avatar
Mario Klingemann committed
140
                if isinstance(y, numbers.Integral):
141
142
                    if y == token_dict_size: #EOS token
                        y = -1
Mario Klingemann's avatar
Mario Klingemann committed
143
                    tokens_temp += [int(y)]
144
                else:
145
146
147
148
149
                    if y.shape[0] == current_embeds.weight.shape[1]:
                        embedding_weights += [y]
                        tokens_temp += [next_new_token]
                        next_new_token += 1
                    else:
150
                        logging.warning("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored {} != {}".format(y.shape[0], current_embeds.weight.shape[1]))
151
            while len(tokens_temp) < len(x):
152
                tokens_temp += [self.special_tokens["pad"]]
153
154
            out_tokens += [tokens_temp]

155
        n = token_dict_size
156
        if len(embedding_weights) > 0:
157
158
            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]
159
160
161
            for x in embedding_weights:
                new_embedding.weight[n] = x
                n += 1
162
            new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
163
            self.transformer.set_input_embeddings(new_embedding)
164
165
166
167
168
169

        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
170

comfyanonymous's avatar
comfyanonymous committed
171
    def forward(self, tokens):
172
        backup_embeds = self.transformer.get_input_embeddings()
173
        device = backup_embeds.weight.device
174
        tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
175
176
        tokens = torch.LongTensor(tokens).to(device)

177
        attention_mask = None
178
        if self.enable_attention_masks or self.zero_out_masked or self.return_attention_masks:
179
            attention_mask = torch.zeros_like(tokens)
180
            end_token = self.special_tokens.get("end", -1)
181
182
183
            for x in range(attention_mask.shape[0]):
                for y in range(attention_mask.shape[1]):
                    attention_mask[x, y] = 1
184
                    if tokens[x, y] == end_token:
185
186
                        break

187
188
189
190
191
        attention_mask_model = None
        if self.enable_attention_masks:
            attention_mask_model = attention_mask

        outputs = self.transformer(tokens, attention_mask_model, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
192
193
194
        self.transformer.set_input_embeddings(backup_embeds)

        if self.layer == "last":
195
            z = outputs[0].float()
comfyanonymous's avatar
comfyanonymous committed
196
        else:
197
198
            z = outputs[1].float()

199
        if self.zero_out_masked:
200
            z *= attention_mask.unsqueeze(-1).float()
201

202
203
204
205
206
207
        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()
208

209
210
211
        if self.return_attention_masks:
            return z, pooled_output, attention_mask

212
        return z, pooled_output
comfyanonymous's avatar
comfyanonymous committed
213
214
215
216

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

217
218
219
    def load_sd(self, sd):
        return self.transformer.load_state_dict(sd, strict=False)

comfyanonymous's avatar
comfyanonymous committed
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
264
265
266
267
268
269
270
271
272
273
274
275
276
277
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

278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
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

297
298
299
300
301
302
303
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)
304

305
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
306
307
308
    if isinstance(embedding_directory, str):
        embedding_directory = [embedding_directory]

309
310
    embedding_directory = expand_directory_list(embedding_directory)

311
312
    valid_file = None
    for embed_dir in embedding_directory:
313
314
315
316
317
318
319
        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
320
321
322
323
324
325
326
        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
327
        else:
328
329
330
331
332
333
334
335
            valid_file = embed_path
        if valid_file is not None:
            break

    if valid_file is None:
        return None

    embed_path = valid_file
336

337
338
    embed_out = None

339
340
341
342
    try:
        if embed_path.lower().endswith(".safetensors"):
            import safetensors.torch
            embed = safetensors.torch.load_file(embed_path, device="cpu")
comfyanonymous's avatar
comfyanonymous committed
343
        else:
344
            if 'weights_only' in torch.load.__code__.co_varnames:
345
346
347
348
                try:
                    embed = torch.load(embed_path, weights_only=True, map_location="cpu")
                except:
                    embed_out = safe_load_embed_zip(embed_path)
349
350
351
            else:
                embed = torch.load(embed_path, map_location="cpu")
    except Exception as e:
352
        logging.warning("{}\n\nerror loading embedding, skipping loading: {}".format(traceback.format_exc(), embedding_name))
353
354
        return None

355
356
357
    if embed_out is None:
        if 'string_to_param' in embed:
            values = embed['string_to_param'].values()
358
359
360
361
362
363
364
365
366
367
            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)
368
369
        elif embed_key is not None and embed_key in embed:
            embed_out = embed[embed_key]
370
371
        else:
            values = embed.values()
372
            embed_out = next(iter(values))
373
    return embed_out
374

375
class SDTokenizer:
376
    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, pad_token=None):
comfyanonymous's avatar
comfyanonymous committed
377
378
        if tokenizer_path is None:
            tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
379
        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
comfyanonymous's avatar
comfyanonymous committed
380
        self.max_length = max_length
381
        self.min_length = min_length
382

comfyanonymous's avatar
comfyanonymous committed
383
        empty = self.tokenizer('')["input_ids"]
384
385
386
387
388
389
390
391
        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]
392
393
394
395
396
397
398
399

        if pad_token is not None:
            self.pad_token = pad_token
        elif pad_with_end:
            self.pad_token = self.end_token
        else:
            self.pad_token = 0

comfyanonymous's avatar
comfyanonymous committed
400
        self.pad_with_end = pad_with_end
401
402
        self.pad_to_max_length = pad_to_max_length

comfyanonymous's avatar
comfyanonymous committed
403
404
        vocab = self.tokenizer.get_vocab()
        self.inv_vocab = {v: k for k, v in vocab.items()}
405
406
        self.embedding_directory = embedding_directory
        self.max_word_length = 8
407
        self.embedding_identifier = "embedding:"
408
        self.embedding_size = embedding_size
409
        self.embedding_key = embedding_key
410

411
    def _try_get_embedding(self, embedding_name:str):
412
413
414
415
        '''
        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.
        '''
416
        embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
417
418
419
        if embed is None:
            stripped = embedding_name.strip(',')
            if len(stripped) < len(embedding_name):
420
                embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
421
422
423
424
                return (embed, embedding_name[len(stripped):])
        return (embed, "")


425
    def tokenize_with_weights(self, text:str, return_word_ids=False):
426
427
428
429
430
431
        '''
        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
        '''
comfyanonymous's avatar
comfyanonymous committed
432
433
434
435

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

436
        #tokenize words
comfyanonymous's avatar
comfyanonymous committed
437
        tokens = []
438
439
440
441
442
443
        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:
444
445
                    embedding_name = word[len(self.embedding_identifier):].strip('\n')
                    embed, leftover = self._try_get_embedding(embedding_name)
446
                    if embed is None:
447
                        logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
448
                    else:
449
                        if len(embed.shape) == 1:
450
                            tokens.append([(embed, weight)])
451
                        else:
452
453
454
455
                            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
456
                    else:
457
458
                        continue
                #parse word
459
                tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
460

461
462
        #reshape token array to CLIP input size
        batched_tokens = []
463
464
465
        batch = []
        if self.start_token is not None:
            batch.append((self.start_token, 1.0, 0))
466
467
        batched_tokens.append(batch)
        for i, t_group in enumerate(tokens):
468
469
            #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
470

471
            while len(t_group) > 0:
BlenderNeko's avatar
BlenderNeko committed
472
473
474
                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
475
476
                    if is_large:
                        batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
BlenderNeko's avatar
BlenderNeko committed
477
                        batch.append((self.end_token, 1.0, 0))
478
                        t_group = t_group[remaining_length:]
BlenderNeko's avatar
BlenderNeko committed
479
                    #add end token and pad
480
                    else:
BlenderNeko's avatar
BlenderNeko committed
481
                        batch.append((self.end_token, 1.0, 0))
482
                        if self.pad_to_max_length:
483
                            batch.extend([(self.pad_token, 1.0, 0)] * (remaining_length))
BlenderNeko's avatar
BlenderNeko committed
484
                    #start new batch
485
486
487
                    batch = []
                    if self.start_token is not None:
                        batch.append((self.start_token, 1.0, 0))
488
                    batched_tokens.append(batch)
489
                else:
490
491
                    batch.extend([(t,w,i+1) for t,w in t_group])
                    t_group = []
492

493
        #fill last batch
494
495
        batch.append((self.end_token, 1.0, 0))
        if self.pad_to_max_length:
496
            batch.extend([(self.pad_token, 1.0, 0)] * (self.max_length - len(batch)))
497
        if self.min_length is not None and len(batch) < self.min_length:
498
            batch.extend([(self.pad_token, 1.0, 0)] * (self.min_length - len(batch)))
comfyanonymous's avatar
comfyanonymous committed
499

500
501
        if not return_word_ids:
            batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
comfyanonymous's avatar
comfyanonymous committed
502

503
        return batched_tokens
comfyanonymous's avatar
comfyanonymous committed
504
505
506
507


    def untokenize(self, token_weight_pair):
        return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525


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):
526
    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, name=None, **kwargs):
527
        super().__init__()
528
529
530
531
532
533
534
535

        if name is not None:
            self.clip_name = name
            self.clip = "{}".format(self.clip_name)
        else:
            self.clip_name = clip_name
            self.clip = "clip_{}".format(self.clip_name)

536
        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
537

538
539
540
541
        self.dtypes = set()
        if dtype is not None:
            self.dtypes.add(dtype)

542
543
    def set_clip_options(self, options):
        getattr(self, self.clip).set_clip_options(options)
544

545
546
    def reset_clip_options(self):
        getattr(self, self.clip).reset_clip_options()
547
548
549
550
551
552
553
554

    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)