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

3
from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextConfig, modeling_utils
4
import comfy.ops
comfyanonymous's avatar
comfyanonymous committed
5
import torch
6
import traceback
7
import zipfile
8
9
from . import model_management
import contextlib
comfyanonymous's avatar
comfyanonymous committed
10

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
56
        if (len(output) == 0):
57
58
            return out[-1:].cpu(), first_pooled
        return torch.cat(output, dim=-2).cpu(), first_pooled
comfyanonymous's avatar
comfyanonymous committed
59

60
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
comfyanonymous's avatar
comfyanonymous committed
61
62
63
64
65
66
67
    """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,
68
69
70
                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, textmodel_path=None, dtype=None,
                 special_tokens={"start": 49406, "end": 49407, "pad": 49407},layer_norm_hidden_state=True, config_class=CLIPTextConfig,
                 model_class=CLIPTextModel, inner_name="text_model"):  # clip-vit-base-patch32
comfyanonymous's avatar
comfyanonymous committed
71
72
        super().__init__()
        assert layer in self.LAYERS
73
        self.num_layers = 12
comfyanonymous's avatar
comfyanonymous committed
74
        if textmodel_path is not None:
75
            self.transformer = model_class.from_pretrained(textmodel_path)
comfyanonymous's avatar
comfyanonymous committed
76
77
78
        else:
            if textmodel_json_config is None:
                textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")
79
            config = config_class.from_json_file(textmodel_json_config)
80
            self.num_layers = config.num_hidden_layers
81
            with comfy.ops.use_comfy_ops(device, dtype):
82
                with modeling_utils.no_init_weights():
83
                    self.transformer = model_class(config)
comfyanonymous's avatar
comfyanonymous committed
84

85
        self.inner_name = inner_name
86
        if dtype is not None:
87
88
            inner_model = getattr(self.transformer, self.inner_name)
            if hasattr(inner_model, "embeddings"):
comfyanonymous's avatar
comfyanonymous committed
89
90
91
92
                embeddings_bak = inner_model.embeddings.to(torch.float32)
                inner_model.embeddings = None
                self.transformer.to(dtype)
                inner_model.embeddings = embeddings_bak
93
            else:
comfyanonymous's avatar
comfyanonymous committed
94
95
96
                previous_inputs = self.transformer.get_input_embeddings().to(torch.float32, copy=True)
                self.transformer.to(dtype)
                self.transformer.set_input_embeddings(previous_inputs)
97

comfyanonymous's avatar
comfyanonymous committed
98
99
100
101
102
        self.max_length = max_length
        if freeze:
            self.freeze()
        self.layer = layer
        self.layer_idx = None
103
        self.special_tokens = special_tokens
104
105
        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))
106
        self.enable_attention_masks = False
107

108
        self.layer_norm_hidden_state = layer_norm_hidden_state
comfyanonymous's avatar
comfyanonymous committed
109
110
        if layer == "hidden":
            assert layer_idx is not None
111
            assert abs(layer_idx) <= self.num_layers
comfyanonymous's avatar
comfyanonymous committed
112
            self.clip_layer(layer_idx)
113
        self.layer_default = (self.layer, self.layer_idx)
comfyanonymous's avatar
comfyanonymous committed
114
115
116
117
118
119
120
121

    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):
122
        if abs(layer_idx) >= self.num_layers:
comfyanonymous's avatar
comfyanonymous committed
123
124
125
126
127
            self.layer = "last"
        else:
            self.layer = "hidden"
            self.layer_idx = layer_idx

128
129
130
131
    def reset_clip_layer(self):
        self.layer = self.layer_default[0]
        self.layer_idx = self.layer_default[1]

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
140
        embedding_weights = []

        for x in tokens:
            tokens_temp = []
            for y in x:
                if isinstance(y, int):
141
142
                    if y == token_dict_size: #EOS token
                        y = -1
143
144
                    tokens_temp += [y]
                else:
145
146
147
148
149
150
                    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])
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
        if getattr(self.transformer, self.inner_name).final_layer_norm.weight.dtype != torch.float32:
178
            precision_scope = torch.autocast
comfyanonymous's avatar
comfyanonymous committed
179
        else:
comfyanonymous's avatar
comfyanonymous committed
180
            precision_scope = lambda a, dtype: contextlib.nullcontext(a)
181

182
        with precision_scope(model_management.get_autocast_device(device), dtype=torch.float32):
183
184
185
186
187
188
189
190
191
192
193
            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(input_ids=tokens, attention_mask=attention_mask, output_hidden_states=self.layer=="hidden")
194
195
196
197
198
199
200
201
202
            self.transformer.set_input_embeddings(backup_embeds)

            if self.layer == "last":
                z = outputs.last_hidden_state
            elif self.layer == "pooled":
                z = outputs.pooler_output[:, None, :]
            else:
                z = outputs.hidden_states[self.layer_idx]
                if self.layer_norm_hidden_state:
203
204
205
206
207
208
                    z = getattr(self.transformer, self.inner_name).final_layer_norm(z)

            if hasattr(outputs, "pooler_output"):
                pooled_output = outputs.pooler_output.float()
            else:
                pooled_output = None
209

210
            if self.text_projection is not None and pooled_output is not None:
211
                pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
212
        return z.float(), pooled_output
comfyanonymous's avatar
comfyanonymous committed
213
214
215
216

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

217
    def load_sd(self, sd):
218
219
220
221
        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)
222
223
        return self.transformer.load_state_dict(sd, strict=False)

comfyanonymous's avatar
comfyanonymous committed
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
278
279
280
281
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

282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
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

301
302
303
304
305
306
307
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)
308

309
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
310
311
312
    if isinstance(embedding_directory, str):
        embedding_directory = [embedding_directory]

313
314
    embedding_directory = expand_directory_list(embedding_directory)

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

    if valid_file is None:
        return None

    embed_path = valid_file
340

341
342
    embed_out = None

343
344
345
346
    try:
        if embed_path.lower().endswith(".safetensors"):
            import safetensors.torch
            embed = safetensors.torch.load_file(embed_path, device="cpu")
comfyanonymous's avatar
comfyanonymous committed
347
        else:
348
            if 'weights_only' in torch.load.__code__.co_varnames:
349
350
351
352
                try:
                    embed = torch.load(embed_path, weights_only=True, map_location="cpu")
                except:
                    embed_out = safe_load_embed_zip(embed_path)
353
354
355
356
357
358
359
360
            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

361
362
363
    if embed_out is None:
        if 'string_to_param' in embed:
            values = embed['string_to_param'].values()
364
365
366
367
368
369
370
371
372
373
            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)
374
375
        elif embed_key is not None and embed_key in embed:
            embed_out = embed[embed_key]
376
377
        else:
            values = embed.values()
378
            embed_out = next(iter(values))
379
    return embed_out
380

381
class SDTokenizer:
382
    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
383
384
        if tokenizer_path is None:
            tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
385
        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
comfyanonymous's avatar
comfyanonymous committed
386
        self.max_length = max_length
387

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

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

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


422
    def tokenize_with_weights(self, text:str, return_word_ids=False):
423
424
425
426
427
428
        '''
        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
429
430
431
432
        if self.pad_with_end:
            pad_token = self.end_token
        else:
            pad_token = 0
comfyanonymous's avatar
comfyanonymous committed
433
434
435
436

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

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

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

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

494
        #fill last batch
495
496
497
        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
498

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

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


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


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):
525
    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
526
527
528
        super().__init__()
        self.clip_name = clip_name
        self.clip = "clip_{}".format(self.clip_name)
529
        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
530
531
532
533
534
535
536
537
538
539
540
541
542
543

    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)