sd1_clip.py 19.8 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

class ClipTokenWeightEncoder:
    def encode_token_weights(self, token_weight_pairs):
13
        to_encode = list(self.empty_tokens)
comfyanonymous's avatar
comfyanonymous committed
14
        for x in token_weight_pairs:
15
16
17
18
19
20
21
22
23
24
25
            tokens = list(map(lambda a: a[0], x))
            to_encode.append(tokens)

        out, pooled = self.encode(to_encode)
        z_empty = out[0:1]
        if pooled.shape[0] > 1:
            first_pooled = pooled[1:2]
        else:
            first_pooled = pooled[0:1]

        output = []
26
27
        for k in range(1, out.shape[0]):
            z = out[k:k+1]
comfyanonymous's avatar
comfyanonymous committed
28
29
            for i in range(len(z)):
                for j in range(len(z[i])):
30
                    weight = token_weight_pairs[k - 1][j][1]
comfyanonymous's avatar
comfyanonymous committed
31
                    z[i][j] = (z[i][j] - z_empty[0][j]) * weight + z_empty[0][j]
32
33
            output.append(z)

comfyanonymous's avatar
comfyanonymous committed
34
        if (len(output) == 0):
35
            return z_empty.cpu(), first_pooled.cpu()
36
        return torch.cat(output, dim=-2).cpu(), first_pooled.cpu()
comfyanonymous's avatar
comfyanonymous committed
37

38
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
comfyanonymous's avatar
comfyanonymous committed
39
40
41
42
43
44
45
    """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,
46
                 freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, textmodel_path=None, dtype=None):  # clip-vit-base-patch32
comfyanonymous's avatar
comfyanonymous committed
47
48
        super().__init__()
        assert layer in self.LAYERS
49
        self.num_layers = 12
comfyanonymous's avatar
comfyanonymous committed
50
51
52
53
54
55
        if textmodel_path is not None:
            self.transformer = CLIPTextModel.from_pretrained(textmodel_path)
        else:
            if textmodel_json_config is None:
                textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")
            config = CLIPTextConfig.from_json_file(textmodel_json_config)
56
            self.num_layers = config.num_hidden_layers
57
            with comfy.ops.use_comfy_ops(device, dtype):
58
59
                with modeling_utils.no_init_weights():
                    self.transformer = CLIPTextModel(config)
comfyanonymous's avatar
comfyanonymous committed
60

61
62
        if dtype is not None:
            self.transformer.to(dtype)
63
64
65
            self.transformer.text_model.embeddings.token_embedding.to(torch.float32)
            self.transformer.text_model.embeddings.position_embedding.to(torch.float32)

comfyanonymous's avatar
comfyanonymous committed
66
67
68
69
70
71
        self.max_length = max_length
        if freeze:
            self.freeze()
        self.layer = layer
        self.layer_idx = None
        self.empty_tokens = [[49406] + [49407] * 76]
72
73
        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))
74
        self.enable_attention_masks = False
75

76
        self.layer_norm_hidden_state = True
comfyanonymous's avatar
comfyanonymous committed
77
78
        if layer == "hidden":
            assert layer_idx is not None
79
            assert abs(layer_idx) <= self.num_layers
comfyanonymous's avatar
comfyanonymous committed
80
            self.clip_layer(layer_idx)
81
        self.layer_default = (self.layer, self.layer_idx)
comfyanonymous's avatar
comfyanonymous committed
82
83
84
85
86
87
88
89

    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):
90
        if abs(layer_idx) >= self.num_layers:
comfyanonymous's avatar
comfyanonymous committed
91
92
93
94
95
            self.layer = "last"
        else:
            self.layer = "hidden"
            self.layer_idx = layer_idx

96
97
98
99
    def reset_clip_layer(self):
        self.layer = self.layer_default[0]
        self.layer_idx = self.layer_default[1]

100
101
    def set_up_textual_embeddings(self, tokens, current_embeds):
        out_tokens = []
102
        next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
103
104
105
106
107
108
        embedding_weights = []

        for x in tokens:
            tokens_temp = []
            for y in x:
                if isinstance(y, int):
109
110
                    if y == token_dict_size: #EOS token
                        y = -1
111
112
                    tokens_temp += [y]
                else:
113
114
115
116
117
118
                    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])
119
120
            while len(tokens_temp) < len(x):
                tokens_temp += [self.empty_tokens[0][-1]]
121
122
            out_tokens += [tokens_temp]

123
        n = token_dict_size
124
        if len(embedding_weights) > 0:
125
126
            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]
127
128
129
            for x in embedding_weights:
                new_embedding.weight[n] = x
                n += 1
130
            new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
131
            self.transformer.set_input_embeddings(new_embedding)
132
133
134
135
136
137

        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
138

comfyanonymous's avatar
comfyanonymous committed
139
    def forward(self, tokens):
140
        backup_embeds = self.transformer.get_input_embeddings()
141
        device = backup_embeds.weight.device
142
        tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
143
144
        tokens = torch.LongTensor(tokens).to(device)

145
        if self.transformer.text_model.final_layer_norm.weight.dtype != torch.float32:
146
            precision_scope = torch.autocast
comfyanonymous's avatar
comfyanonymous committed
147
        else:
148
            precision_scope = lambda a, b: contextlib.nullcontext(a)
149

150
        with precision_scope(model_management.get_autocast_device(device), torch.float32):
151
152
153
154
155
156
157
158
159
160
161
            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")
162
163
164
165
166
167
168
169
170
171
172
173
174
            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:
                    z = self.transformer.text_model.final_layer_norm(z)

            pooled_output = outputs.pooler_output
            if self.text_projection is not None:
175
                pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
176
        return z.float(), pooled_output.float()
comfyanonymous's avatar
comfyanonymous committed
177
178
179
180

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

181
    def load_sd(self, sd):
182
183
184
185
        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)
186
187
        return self.transformer.load_state_dict(sd, strict=False)

comfyanonymous's avatar
comfyanonymous committed
188
189
190
191
192
193
194
195
196
197
198
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
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

246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
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

265
266
267
268
269
270
271
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)
272

273
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
274
275
276
    if isinstance(embedding_directory, str):
        embedding_directory = [embedding_directory]

277
278
    embedding_directory = expand_directory_list(embedding_directory)

279
280
    valid_file = None
    for embed_dir in embedding_directory:
281
282
283
284
285
286
287
        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
288
289
290
291
292
293
294
        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
295
        else:
296
297
298
299
300
301
302
303
            valid_file = embed_path
        if valid_file is not None:
            break

    if valid_file is None:
        return None

    embed_path = valid_file
304

305
306
    embed_out = None

307
308
309
310
    try:
        if embed_path.lower().endswith(".safetensors"):
            import safetensors.torch
            embed = safetensors.torch.load_file(embed_path, device="cpu")
comfyanonymous's avatar
comfyanonymous committed
311
        else:
312
            if 'weights_only' in torch.load.__code__.co_varnames:
313
314
315
316
                try:
                    embed = torch.load(embed_path, weights_only=True, map_location="cpu")
                except:
                    embed_out = safe_load_embed_zip(embed_path)
317
318
319
320
321
322
323
324
            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

325
326
327
    if embed_out is None:
        if 'string_to_param' in embed:
            values = embed['string_to_param'].values()
328
329
330
331
332
333
334
335
336
337
            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)
338
339
        elif embed_key is not None and embed_key in embed:
            embed_out = embed[embed_key]
340
341
        else:
            values = embed.values()
342
            embed_out = next(iter(values))
343
    return embed_out
344

345
class SDTokenizer:
346
    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
347
348
        if tokenizer_path is None:
            tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
349
        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
comfyanonymous's avatar
comfyanonymous committed
350
        self.max_length = max_length
351

comfyanonymous's avatar
comfyanonymous committed
352
        empty = self.tokenizer('')["input_ids"]
353
354
355
356
357
358
359
360
        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
361
        self.pad_with_end = pad_with_end
362
363
        self.pad_to_max_length = pad_to_max_length

comfyanonymous's avatar
comfyanonymous committed
364
365
        vocab = self.tokenizer.get_vocab()
        self.inv_vocab = {v: k for k, v in vocab.items()}
366
367
        self.embedding_directory = embedding_directory
        self.max_word_length = 8
368
        self.embedding_identifier = "embedding:"
369
        self.embedding_size = embedding_size
370
        self.embedding_key = embedding_key
371

372
    def _try_get_embedding(self, embedding_name:str):
373
374
375
376
        '''
        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.
        '''
377
        embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
378
379
380
        if embed is None:
            stripped = embedding_name.strip(',')
            if len(stripped) < len(embedding_name):
381
                embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
382
383
384
385
                return (embed, embedding_name[len(stripped):])
        return (embed, "")


386
    def tokenize_with_weights(self, text:str, return_word_ids=False):
387
388
389
390
391
392
        '''
        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
393
394
395
396
        if self.pad_with_end:
            pad_token = self.end_token
        else:
            pad_token = 0
comfyanonymous's avatar
comfyanonymous committed
397
398
399
400

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

401
        #tokenize words
comfyanonymous's avatar
comfyanonymous committed
402
        tokens = []
403
404
405
406
407
408
        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:
409
410
                    embedding_name = word[len(self.embedding_identifier):].strip('\n')
                    embed, leftover = self._try_get_embedding(embedding_name)
411
                    if embed is None:
412
                        print(f"warning, embedding:{embedding_name} does not exist, ignoring")
413
                    else:
414
                        if len(embed.shape) == 1:
415
                            tokens.append([(embed, weight)])
416
                        else:
417
418
419
420
                            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
421
                    else:
422
423
                        continue
                #parse word
424
                tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
425

426
427
        #reshape token array to CLIP input size
        batched_tokens = []
428
429
430
        batch = []
        if self.start_token is not None:
            batch.append((self.start_token, 1.0, 0))
431
432
        batched_tokens.append(batch)
        for i, t_group in enumerate(tokens):
433
434
            #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
435

436
            while len(t_group) > 0:
BlenderNeko's avatar
BlenderNeko committed
437
438
439
                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
440
441
                    if is_large:
                        batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
BlenderNeko's avatar
BlenderNeko committed
442
                        batch.append((self.end_token, 1.0, 0))
443
                        t_group = t_group[remaining_length:]
BlenderNeko's avatar
BlenderNeko committed
444
                    #add end token and pad
445
                    else:
BlenderNeko's avatar
BlenderNeko committed
446
                        batch.append((self.end_token, 1.0, 0))
447
448
                        if self.pad_to_max_length:
                            batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
BlenderNeko's avatar
BlenderNeko committed
449
                    #start new batch
450
451
452
                    batch = []
                    if self.start_token is not None:
                        batch.append((self.start_token, 1.0, 0))
453
                    batched_tokens.append(batch)
454
                else:
455
456
                    batch.extend([(t,w,i+1) for t,w in t_group])
                    t_group = []
457

458
        #fill last batch
459
460
461
        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
462

463
464
        if not return_word_ids:
            batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
comfyanonymous's avatar
comfyanonymous committed
465

466
        return batched_tokens
comfyanonymous's avatar
comfyanonymous committed
467
468
469
470


    def untokenize(self, token_weight_pair):
        return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488


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):
489
    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
490
491
492
        super().__init__()
        self.clip_name = clip_name
        self.clip = "clip_{}".format(self.clip_name)
493
        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
494
495
496
497
498
499
500
501
502
503
504
505
506
507

    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)