sd1_clip.py 19.9 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
comfyanonymous's avatar
comfyanonymous committed
11

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
57
        if (len(output) == 0):
58
59
            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
60

61
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
comfyanonymous's avatar
comfyanonymous committed
62
63
64
65
66
67
68
    """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,
69
70
                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=comfy.clip_model.CLIPTextModel,
                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True):  # clip-vit-base-patch32
comfyanonymous's avatar
comfyanonymous committed
71
72
        super().__init__()
        assert layer in self.LAYERS
73
74
75
76
77
78
79

        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)

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

comfyanonymous's avatar
comfyanonymous committed
83
84
85
86
87
        self.max_length = max_length
        if freeze:
            self.freeze()
        self.layer = layer
        self.layer_idx = None
88
        self.special_tokens = special_tokens
89
90
        self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
        self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
91
        self.enable_attention_masks = False
92

93
        self.layer_norm_hidden_state = layer_norm_hidden_state
comfyanonymous's avatar
comfyanonymous committed
94
95
        if layer == "hidden":
            assert layer_idx is not None
96
            assert abs(layer_idx) < self.num_layers
comfyanonymous's avatar
comfyanonymous committed
97
            self.clip_layer(layer_idx)
98
        self.layer_default = (self.layer, self.layer_idx)
comfyanonymous's avatar
comfyanonymous committed
99
100
101
102
103
104
105
106

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

    def clip_layer(self, layer_idx):
107
        if abs(layer_idx) > self.num_layers:
comfyanonymous's avatar
comfyanonymous committed
108
109
110
111
112
            self.layer = "last"
        else:
            self.layer = "hidden"
            self.layer_idx = layer_idx

113
114
115
116
    def reset_clip_layer(self):
        self.layer = self.layer_default[0]
        self.layer_idx = self.layer_default[1]

117
118
    def set_up_textual_embeddings(self, tokens, current_embeds):
        out_tokens = []
119
        next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
120
121
122
123
124
125
        embedding_weights = []

        for x in tokens:
            tokens_temp = []
            for y in x:
                if isinstance(y, int):
126
127
                    if y == token_dict_size: #EOS token
                        y = -1
128
129
                    tokens_temp += [y]
                else:
130
131
132
133
134
135
                    if y.shape[0] == current_embeds.weight.shape[1]:
                        embedding_weights += [y]
                        tokens_temp += [next_new_token]
                        next_new_token += 1
                    else:
                        print("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored", y.shape[0], current_embeds.weight.shape[1])
136
            while len(tokens_temp) < len(x):
137
                tokens_temp += [self.special_tokens["pad"]]
138
139
            out_tokens += [tokens_temp]

140
        n = token_dict_size
141
        if len(embedding_weights) > 0:
142
143
            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]
144
145
146
            for x in embedding_weights:
                new_embedding.weight[n] = x
                n += 1
147
            new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
148
            self.transformer.set_input_embeddings(new_embedding)
149
150
151
152
153
154

        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
155

comfyanonymous's avatar
comfyanonymous committed
156
    def forward(self, tokens):
157
        backup_embeds = self.transformer.get_input_embeddings()
158
        device = backup_embeds.weight.device
159
        tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
160
161
        tokens = torch.LongTensor(tokens).to(device)

162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
        attention_mask = None
        if self.enable_attention_masks:
            attention_mask = torch.zeros_like(tokens)
            max_token = self.transformer.get_input_embeddings().weight.shape[0] - 1
            for x in range(attention_mask.shape[0]):
                for y in range(attention_mask.shape[1]):
                    attention_mask[x, y] = 1
                    if tokens[x, y] == max_token:
                        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":
            z = outputs[0]
comfyanonymous's avatar
comfyanonymous committed
177
        else:
178
179
180
181
182
183
184
185
186
            z = outputs[1]

        if outputs[2] is not None:
            pooled_output = outputs[2].float()
        else:
            pooled_output = None

        if self.text_projection is not None and pooled_output is not None:
            pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
187
        return z.float(), pooled_output
comfyanonymous's avatar
comfyanonymous committed
188
189
190
191

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

192
    def load_sd(self, sd):
193
194
195
196
        if "text_projection" in sd:
            self.text_projection[:] = sd.pop("text_projection")
        if "text_projection.weight" in sd:
            self.text_projection[:] = sd.pop("text_projection.weight").transpose(0, 1)
197
198
        return self.transformer.load_state_dict(sd, strict=False)

comfyanonymous's avatar
comfyanonymous committed
199
200
201
202
203
204
205
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
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

257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
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

276
277
278
279
280
281
282
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)
283

284
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
285
286
287
    if isinstance(embedding_directory, str):
        embedding_directory = [embedding_directory]

288
289
    embedding_directory = expand_directory_list(embedding_directory)

290
291
    valid_file = None
    for embed_dir in embedding_directory:
292
293
294
295
296
297
298
        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
299
300
301
302
303
304
305
        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
306
        else:
307
308
309
310
311
312
313
314
            valid_file = embed_path
        if valid_file is not None:
            break

    if valid_file is None:
        return None

    embed_path = valid_file
315

316
317
    embed_out = None

318
319
320
321
    try:
        if embed_path.lower().endswith(".safetensors"):
            import safetensors.torch
            embed = safetensors.torch.load_file(embed_path, device="cpu")
comfyanonymous's avatar
comfyanonymous committed
322
        else:
323
            if 'weights_only' in torch.load.__code__.co_varnames:
324
325
326
327
                try:
                    embed = torch.load(embed_path, weights_only=True, map_location="cpu")
                except:
                    embed_out = safe_load_embed_zip(embed_path)
328
329
330
331
332
333
334
335
            else:
                embed = torch.load(embed_path, map_location="cpu")
    except Exception as e:
        print(traceback.format_exc())
        print()
        print("error loading embedding, skipping loading:", embedding_name)
        return None

336
337
338
    if embed_out is None:
        if 'string_to_param' in embed:
            values = embed['string_to_param'].values()
339
340
341
342
343
344
345
346
347
348
            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)
349
350
        elif embed_key is not None and embed_key in embed:
            embed_out = embed[embed_key]
351
352
        else:
            values = embed.values()
353
            embed_out = next(iter(values))
354
    return embed_out
355

356
class SDTokenizer:
357
    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):
comfyanonymous's avatar
comfyanonymous committed
358
359
        if tokenizer_path is None:
            tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
360
        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
comfyanonymous's avatar
comfyanonymous committed
361
        self.max_length = max_length
362

comfyanonymous's avatar
comfyanonymous committed
363
        empty = self.tokenizer('')["input_ids"]
364
365
366
367
368
369
370
371
        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
372
        self.pad_with_end = pad_with_end
373
374
        self.pad_to_max_length = pad_to_max_length

comfyanonymous's avatar
comfyanonymous committed
375
376
        vocab = self.tokenizer.get_vocab()
        self.inv_vocab = {v: k for k, v in vocab.items()}
377
378
        self.embedding_directory = embedding_directory
        self.max_word_length = 8
379
        self.embedding_identifier = "embedding:"
380
        self.embedding_size = embedding_size
381
        self.embedding_key = embedding_key
382

383
    def _try_get_embedding(self, embedding_name:str):
384
385
386
387
        '''
        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.
        '''
388
        embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
389
390
391
        if embed is None:
            stripped = embedding_name.strip(',')
            if len(stripped) < len(embedding_name):
392
                embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
393
394
395
396
                return (embed, embedding_name[len(stripped):])
        return (embed, "")


397
    def tokenize_with_weights(self, text:str, return_word_ids=False):
398
399
400
401
402
403
        '''
        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
404
405
406
407
        if self.pad_with_end:
            pad_token = self.end_token
        else:
            pad_token = 0
comfyanonymous's avatar
comfyanonymous committed
408
409
410
411

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

412
        #tokenize words
comfyanonymous's avatar
comfyanonymous committed
413
        tokens = []
414
415
416
417
418
419
        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:
420
421
                    embedding_name = word[len(self.embedding_identifier):].strip('\n')
                    embed, leftover = self._try_get_embedding(embedding_name)
422
                    if embed is None:
423
                        print(f"warning, embedding:{embedding_name} does not exist, ignoring")
424
                    else:
425
                        if len(embed.shape) == 1:
426
                            tokens.append([(embed, weight)])
427
                        else:
428
429
430
431
                            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
432
                    else:
433
434
                        continue
                #parse word
435
                tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
436

437
438
        #reshape token array to CLIP input size
        batched_tokens = []
439
440
441
        batch = []
        if self.start_token is not None:
            batch.append((self.start_token, 1.0, 0))
442
443
        batched_tokens.append(batch)
        for i, t_group in enumerate(tokens):
444
445
            #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
446

447
            while len(t_group) > 0:
BlenderNeko's avatar
BlenderNeko committed
448
449
450
                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
451
452
                    if is_large:
                        batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
BlenderNeko's avatar
BlenderNeko committed
453
                        batch.append((self.end_token, 1.0, 0))
454
                        t_group = t_group[remaining_length:]
BlenderNeko's avatar
BlenderNeko committed
455
                    #add end token and pad
456
                    else:
BlenderNeko's avatar
BlenderNeko committed
457
                        batch.append((self.end_token, 1.0, 0))
458
459
                        if self.pad_to_max_length:
                            batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
BlenderNeko's avatar
BlenderNeko committed
460
                    #start new batch
461
462
463
                    batch = []
                    if self.start_token is not None:
                        batch.append((self.start_token, 1.0, 0))
464
                    batched_tokens.append(batch)
465
                else:
466
467
                    batch.extend([(t,w,i+1) for t,w in t_group])
                    t_group = []
468

469
        #fill last batch
470
471
472
        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)))
comfyanonymous's avatar
comfyanonymous committed
473

474
475
        if not return_word_ids:
            batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
comfyanonymous's avatar
comfyanonymous committed
476

477
        return batched_tokens
comfyanonymous's avatar
comfyanonymous committed
478
479
480
481


    def untokenize(self, token_weight_pair):
        return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499


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):
500
    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
501
502
503
        super().__init__()
        self.clip_name = clip_name
        self.clip = "clip_{}".format(self.clip_name)
504
        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
505
506
507
508
509
510
511
512
513
514
515
516
517
518

    def clip_layer(self, layer_idx):
        getattr(self, self.clip).clip_layer(layer_idx)

    def reset_clip_layer(self):
        getattr(self, self.clip).reset_clip_layer()

    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)