"src/vscode:/vscode.git/clone" did not exist on "26df5406e4c72f95b3825df633c25efe30913224"
sd.py 56.9 KB
Newer Older
comfyanonymous's avatar
comfyanonymous committed
1
import torch
2
import contextlib
3
import copy
4
import inspect
5
import math
comfyanonymous's avatar
comfyanonymous committed
6

7
from comfy import model_management
8
9
from .ldm.util import instantiate_from_config
from .ldm.models.autoencoder import AutoencoderKL
10
import yaml
comfyanonymous's avatar
comfyanonymous committed
11
from .cldm import cldm
12
from .t2i_adapter import adapter
comfyanonymous's avatar
comfyanonymous committed
13
14

from . import utils
15
from . import clip_vision
16
from . import gligen
17
from . import diffusers_convert
comfyanonymous's avatar
comfyanonymous committed
18
from . import model_base
19
from . import model_detection
20

21
22
from . import sd1_clip
from . import sd2_clip
23
from . import sdxl_clip
comfyanonymous's avatar
comfyanonymous committed
24

25
def load_model_weights(model, sd):
comfyanonymous's avatar
comfyanonymous committed
26
    m, u = model.load_state_dict(sd, strict=False)
27
28
    m = set(m)
    unexpected_keys = set(u)
comfyanonymous's avatar
comfyanonymous committed
29
30
31

    k = list(sd.keys())
    for x in k:
32
33
34
35
36
37
38
39
40
41
        if x not in unexpected_keys:
            w = sd.pop(x)
            del w
    if len(m) > 0:
        print("missing", m)
    return model

def load_clip_weights(model, sd):
    k = list(sd.keys())
    for x in k:
comfyanonymous's avatar
comfyanonymous committed
42
43
44
45
        if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):
            y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")
            sd[y] = sd.pop(x)

comfyanonymous's avatar
comfyanonymous committed
46
47
48
49
    if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in sd:
        ids = sd['cond_stage_model.transformer.text_model.embeddings.position_ids']
        if ids.dtype == torch.float32:
            sd['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
50

51
52
    sd = utils.transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.text_model.", 24)
    return load_model_weights(model, sd)
comfyanonymous's avatar
comfyanonymous committed
53

54
55
56
57
58
59
60
61
62
63
LORA_CLIP_MAP = {
    "mlp.fc1": "mlp_fc1",
    "mlp.fc2": "mlp_fc2",
    "self_attn.k_proj": "self_attn_k_proj",
    "self_attn.q_proj": "self_attn_q_proj",
    "self_attn.v_proj": "self_attn_v_proj",
    "self_attn.out_proj": "self_attn_out_proj",
}


64
def load_lora(lora, to_load):
65
66
67
    patch_dict = {}
    loaded_keys = set()
    for x in to_load:
comfyanonymous's avatar
comfyanonymous committed
68
69
70
71
72
73
        alpha_name = "{}.alpha".format(x)
        alpha = None
        if alpha_name in lora.keys():
            alpha = lora[alpha_name].item()
            loaded_keys.add(alpha_name)

74
75
        regular_lora = "{}.lora_up.weight".format(x)
        diffusers_lora = "{}_lora.up.weight".format(x)
76
        transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
77
78
79
80
81
82
83
84
85
86
        A_name = None

        if regular_lora in lora.keys():
            A_name = regular_lora
            B_name = "{}.lora_down.weight".format(x)
            mid_name = "{}.lora_mid.weight".format(x)
        elif diffusers_lora in lora.keys():
            A_name = diffusers_lora
            B_name = "{}_lora.down.weight".format(x)
            mid_name = None
87
88
89
90
        elif transformers_lora in lora.keys():
            A_name = transformers_lora
            B_name ="{}.lora_linear_layer.down.weight".format(x)
            mid_name = None
91
92

        if A_name is not None:
93
            mid = None
94
            if mid_name is not None and mid_name in lora.keys():
95
96
97
                mid = lora[mid_name]
                loaded_keys.add(mid_name)
            patch_dict[to_load[x]] = (lora[A_name], lora[B_name], alpha, mid)
98
99
            loaded_keys.add(A_name)
            loaded_keys.add(B_name)
comfyanonymous's avatar
comfyanonymous committed
100

comfyanonymous's avatar
comfyanonymous committed
101
102

        ######## loha
comfyanonymous's avatar
comfyanonymous committed
103
104
105
106
        hada_w1_a_name = "{}.hada_w1_a".format(x)
        hada_w1_b_name = "{}.hada_w1_b".format(x)
        hada_w2_a_name = "{}.hada_w2_a".format(x)
        hada_w2_b_name = "{}.hada_w2_b".format(x)
107
108
        hada_t1_name = "{}.hada_t1".format(x)
        hada_t2_name = "{}.hada_t2".format(x)
comfyanonymous's avatar
comfyanonymous committed
109
        if hada_w1_a_name in lora.keys():
110
111
112
113
114
115
116
117
118
            hada_t1 = None
            hada_t2 = None
            if hada_t1_name in lora.keys():
                hada_t1 = lora[hada_t1_name]
                hada_t2 = lora[hada_t2_name]
                loaded_keys.add(hada_t1_name)
                loaded_keys.add(hada_t2_name)

            patch_dict[to_load[x]] = (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2)
comfyanonymous's avatar
comfyanonymous committed
119
120
121
122
123
            loaded_keys.add(hada_w1_a_name)
            loaded_keys.add(hada_w1_b_name)
            loaded_keys.add(hada_w2_a_name)
            loaded_keys.add(hada_w2_b_name)

comfyanonymous's avatar
comfyanonymous committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

        ######## lokr
        lokr_w1_name = "{}.lokr_w1".format(x)
        lokr_w2_name = "{}.lokr_w2".format(x)
        lokr_w1_a_name = "{}.lokr_w1_a".format(x)
        lokr_w1_b_name = "{}.lokr_w1_b".format(x)
        lokr_t2_name = "{}.lokr_t2".format(x)
        lokr_w2_a_name = "{}.lokr_w2_a".format(x)
        lokr_w2_b_name = "{}.lokr_w2_b".format(x)

        lokr_w1 = None
        if lokr_w1_name in lora.keys():
            lokr_w1 = lora[lokr_w1_name]
            loaded_keys.add(lokr_w1_name)

        lokr_w2 = None
        if lokr_w2_name in lora.keys():
            lokr_w2 = lora[lokr_w2_name]
            loaded_keys.add(lokr_w2_name)

        lokr_w1_a = None
        if lokr_w1_a_name in lora.keys():
            lokr_w1_a = lora[lokr_w1_a_name]
            loaded_keys.add(lokr_w1_a_name)

        lokr_w1_b = None
        if lokr_w1_b_name in lora.keys():
            lokr_w1_b = lora[lokr_w1_b_name]
            loaded_keys.add(lokr_w1_b_name)

        lokr_w2_a = None
        if lokr_w2_a_name in lora.keys():
            lokr_w2_a = lora[lokr_w2_a_name]
            loaded_keys.add(lokr_w2_a_name)

        lokr_w2_b = None
        if lokr_w2_b_name in lora.keys():
            lokr_w2_b = lora[lokr_w2_b_name]
            loaded_keys.add(lokr_w2_b_name)

        lokr_t2 = None
        if lokr_t2_name in lora.keys():
            lokr_t2 = lora[lokr_t2_name]
            loaded_keys.add(lokr_t2_name)

        if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
            patch_dict[to_load[x]] = (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2)

172
173
174
175
176
    for x in lora.keys():
        if x not in loaded_keys:
            print("lora key not loaded", x)
    return patch_dict

177
def model_lora_keys_clip(model, key_map={}):
178
179
    sdk = model.state_dict().keys()

comfyanonymous's avatar
comfyanonymous committed
180
    text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}"
181
182
    clip_l_present = False
    for b in range(32):
183
184
185
        for c in LORA_CLIP_MAP:
            k = "transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
            if k in sdk:
comfyanonymous's avatar
comfyanonymous committed
186
                lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])
187
                key_map[lora_key] = k
comfyanonymous's avatar
comfyanonymous committed
188
189
                lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c])
                key_map[lora_key] = k
190
191
                lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
                key_map[lora_key] = k
comfyanonymous's avatar
comfyanonymous committed
192

193
194
195
196
197
            k = "clip_l.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
            if k in sdk:
                lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
                key_map[lora_key] = k
                clip_l_present = True
198
199
                lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
                key_map[lora_key] = k
200
201
202
203
204

            k = "clip_g.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
            if k in sdk:
                if clip_l_present:
                    lora_key = "lora_te2_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
205
206
207
                    key_map[lora_key] = k
                    lora_key = "text_encoder_2.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
                    key_map[lora_key] = k
208
209
                else:
                    lora_key = "lora_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #TODO: test if this is correct for SDXL-Refiner
210
211
212
                    key_map[lora_key] = k
                    lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
                    key_map[lora_key] = k
213

214
    return key_map
comfyanonymous's avatar
comfyanonymous committed
215

216
217
def model_lora_keys_unet(model, key_map={}):
    sdk = model.state_dict().keys()
comfyanonymous's avatar
comfyanonymous committed
218

219
220
221
222
223
    for k in sdk:
        if k.startswith("diffusion_model.") and k.endswith(".weight"):
            key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
            key_map["lora_unet_{}".format(key_lora)] = k

224
225
226
    diffusers_keys = utils.unet_to_diffusers(model.model_config.unet_config)
    for k in diffusers_keys:
        if k.endswith(".weight"):
227
            unet_key = "diffusion_model.{}".format(diffusers_keys[k])
228
            key_lora = k[:-len(".weight")].replace(".", "_")
229
230
231
232
233
234
235
236
            key_map["lora_unet_{}".format(key_lora)] = unet_key

            diffusers_lora_prefix = ["", "unet."]
            for p in diffusers_lora_prefix:
                diffusers_lora_key = "{}{}".format(p, k[:-len(".weight")].replace(".to_", ".processor.to_"))
                if diffusers_lora_key.endswith(".to_out.0"):
                    diffusers_lora_key = diffusers_lora_key[:-2]
                key_map[diffusers_lora_key] = unet_key
237
238
    return key_map

239
240
241
242
243
244
245
246
def set_attr(obj, attr, value):
    attrs = attr.split(".")
    for name in attrs[:-1]:
        obj = getattr(obj, name)
    prev = getattr(obj, attrs[-1])
    setattr(obj, attrs[-1], torch.nn.Parameter(value))
    del prev

comfyanonymous's avatar
comfyanonymous committed
247
248
249
250
251
252
253
def get_attr(obj, attr):
    attrs = attr.split(".")
    for name in attrs:
        obj = getattr(obj, name)
    return obj


254
class ModelPatcher:
comfyanonymous's avatar
comfyanonymous committed
255
    def __init__(self, model, load_device, offload_device, size=0, current_device=None):
256
        self.size = size
257
        self.model = model
258
        self.patches = {}
259
        self.backup = {}
260
        self.model_options = {"transformer_options":{}}
261
        self.model_size()
262
263
        self.load_device = load_device
        self.offload_device = offload_device
comfyanonymous's avatar
comfyanonymous committed
264
265
266
267
        if current_device is None:
            self.current_device = self.offload_device
        else:
            self.current_device = current_device
268
269
270
271
272
273
274
275
276
277

    def model_size(self):
        if self.size > 0:
            return self.size
        model_sd = self.model.state_dict()
        size = 0
        for k in model_sd:
            t = model_sd[k]
            size += t.nelement() * t.element_size()
        self.size = size
278
        self.model_keys = set(model_sd.keys())
279
        return size
280
281

    def clone(self):
comfyanonymous's avatar
comfyanonymous committed
282
        n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device)
283
284
285
286
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]

287
        n.model_options = copy.deepcopy(self.model_options)
288
        n.model_keys = self.model_keys
289
290
        return n

comfyanonymous's avatar
comfyanonymous committed
291
292
293
294
295
    def is_clone(self, other):
        if hasattr(other, 'model') and self.model is other.model:
            return True
        return False

296
    def set_model_sampler_cfg_function(self, sampler_cfg_function):
297
298
299
300
        if len(inspect.signature(sampler_cfg_function).parameters) == 3:
            self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
        else:
            self.model_options["sampler_cfg_function"] = sampler_cfg_function
301

302
303
304
    def set_model_unet_function_wrapper(self, unet_wrapper_function):
        self.model_options["model_function_wrapper"] = unet_wrapper_function

305
306
307
308
309
310
    def set_model_patch(self, patch, name):
        to = self.model_options["transformer_options"]
        if "patches" not in to:
            to["patches"] = {}
        to["patches"][name] = to["patches"].get(name, []) + [patch]

311
312
313
314
315
316
317
318
    def set_model_patch_replace(self, patch, name, block_name, number):
        to = self.model_options["transformer_options"]
        if "patches_replace" not in to:
            to["patches_replace"] = {}
        if name not in to["patches_replace"]:
            to["patches_replace"][name] = {}
        to["patches_replace"][name][(block_name, number)] = patch

319
320
321
322
323
324
    def set_model_attn1_patch(self, patch):
        self.set_model_patch(patch, "attn1_patch")

    def set_model_attn2_patch(self, patch):
        self.set_model_patch(patch, "attn2_patch")

325
326
327
328
329
330
331
332
333
    def set_model_attn1_replace(self, patch, block_name, number):
        self.set_model_patch_replace(patch, "attn1", block_name, number)

    def set_model_attn2_replace(self, patch, block_name, number):
        self.set_model_patch_replace(patch, "attn2", block_name, number)

    def set_model_attn1_output_patch(self, patch):
        self.set_model_patch(patch, "attn1_output_patch")

334
335
336
    def set_model_attn2_output_patch(self, patch):
        self.set_model_patch(patch, "attn2_output_patch")

337
338
339
340
341
342
343
344
345
    def model_patches_to(self, device):
        to = self.model_options["transformer_options"]
        if "patches" in to:
            patches = to["patches"]
            for name in patches:
                patch_list = patches[name]
                for i in range(len(patch_list)):
                    if hasattr(patch_list[i], "to"):
                        patch_list[i] = patch_list[i].to(device)
346
347
348
349
350
351
352
        if "patches_replace" in to:
            patches = to["patches_replace"]
            for name in patches:
                patch_list = patches[name]
                for k in patch_list:
                    if hasattr(patch_list[k], "to"):
                        patch_list[k] = patch_list[k].to(device)
353

354
    def model_dtype(self):
355
356
        if hasattr(self.model, "get_dtype"):
            return self.model.get_dtype()
357

358
    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
359
        p = set()
360
        for k in patches:
361
            if k in self.model_keys:
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
                p.add(k)
                current_patches = self.patches.get(k, [])
                current_patches.append((strength_patch, patches[k], strength_model))
                self.patches[k] = current_patches

        return list(p)

    def get_key_patches(self, filter_prefix=None):
        model_sd = self.model_state_dict()
        p = {}
        for k in model_sd:
            if filter_prefix is not None:
                if not k.startswith(filter_prefix):
                    continue
            if k in self.patches:
                p[k] = [model_sd[k]] + self.patches[k]
            else:
                p[k] = (model_sd[k],)
        return p
381

382
    def model_state_dict(self, filter_prefix=None):
383
384
        sd = self.model.state_dict()
        keys = list(sd.keys())
385
386
387
388
        if filter_prefix is not None:
            for k in keys:
                if not k.startswith(filter_prefix):
                    sd.pop(k)
389
390
        return sd

391
    def patch_model(self, device_to=None):
392
        model_sd = self.model_state_dict()
393
394
395
396
        for key in self.patches:
            if key not in model_sd:
                print("could not patch. key doesn't exist in model:", k)
                continue
397

398
            weight = model_sd[key]
399

400
            if key not in self.backup:
401
                self.backup[key] = weight.to(self.offload_device)
402

403
404
405
406
            if device_to is not None:
                temp_weight = weight.float().to(device_to, copy=True)
            else:
                temp_weight = weight.to(torch.float32, copy=True)
407
408
            out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
            set_attr(self.model, key, out_weight)
409
            del temp_weight
comfyanonymous's avatar
comfyanonymous committed
410
411
412
413
414

        if device_to is not None:
            self.model.to(device_to)
            self.current_device = device_to

415
        return self.model
comfyanonymous's avatar
comfyanonymous committed
416

417
418
419
420
421
422
423
424
425
426
427
428
429
430
    def calculate_weight(self, patches, weight, key):
        for p in patches:
            alpha = p[0]
            v = p[1]
            strength_model = p[2]

            if strength_model != 1.0:
                weight *= strength_model

            if isinstance(v, list):
                v = (self.calculate_weight(v[1:], v[0].clone(), key), )

            if len(v) == 1:
                w1 = v[0]
431
432
433
434
435
                if alpha != 0.0:
                    if w1.shape != weight.shape:
                        print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
                    else:
                        weight += alpha * w1.type(weight.dtype).to(weight.device)
436
            elif len(v) == 4: #lora/locon
comfyanonymous's avatar
comfyanonymous committed
437
438
                mat1 = v[0].float().to(weight.device)
                mat2 = v[1].float().to(weight.device)
439
440
441
442
                if v[2] is not None:
                    alpha *= v[2] / mat2.shape[0]
                if v[3] is not None:
                    #locon mid weights, hopefully the math is fine because I didn't properly test it
comfyanonymous's avatar
comfyanonymous committed
443
444
445
                    mat3 = v[3].float().to(weight.device)
                    final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
                    mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
446
447
448
449
                try:
                    weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
450
451
452
453
454
455
456
457
458
459
460
461
462
            elif len(v) == 8: #lokr
                w1 = v[0]
                w2 = v[1]
                w1_a = v[3]
                w1_b = v[4]
                w2_a = v[5]
                w2_b = v[6]
                t2 = v[7]
                dim = None

                if w1 is None:
                    dim = w1_b.shape[0]
                    w1 = torch.mm(w1_a.float(), w1_b.float())
comfyanonymous's avatar
comfyanonymous committed
463
464
                else:
                    w1 = w1.float().to(weight.device)
465
466
467
468

                if w2 is None:
                    dim = w2_b.shape[0]
                    if t2 is None:
comfyanonymous's avatar
comfyanonymous committed
469
                        w2 = torch.mm(w2_a.float().to(weight.device), w2_b.float().to(weight.device))
470
                    else:
comfyanonymous's avatar
comfyanonymous committed
471
472
473
                        w2 = torch.einsum('i j k l, j r, i p -> p r k l', t2.float().to(weight.device), w2_b.float().to(weight.device), w2_a.float().to(weight.device))
                else:
                    w2 = w2.float().to(weight.device)
474
475
476
477
478
479

                if len(w2.shape) == 4:
                    w1 = w1.unsqueeze(2).unsqueeze(2)
                if v[2] is not None and dim is not None:
                    alpha *= v[2] / dim

480
481
482
483
                try:
                    weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
484
485
486
487
488
489
490
491
492
493
            else: #loha
                w1a = v[0]
                w1b = v[1]
                if v[2] is not None:
                    alpha *= v[2] / w1b.shape[0]
                w2a = v[3]
                w2b = v[4]
                if v[5] is not None: #cp decomposition
                    t1 = v[5]
                    t2 = v[6]
comfyanonymous's avatar
comfyanonymous committed
494
495
                    m1 = torch.einsum('i j k l, j r, i p -> p r k l', t1.float().to(weight.device), w1b.float().to(weight.device), w1a.float().to(weight.device))
                    m2 = torch.einsum('i j k l, j r, i p -> p r k l', t2.float().to(weight.device), w2b.float().to(weight.device), w2a.float().to(weight.device))
496
                else:
comfyanonymous's avatar
comfyanonymous committed
497
498
                    m1 = torch.mm(w1a.float().to(weight.device), w1b.float().to(weight.device))
                    m2 = torch.mm(w2a.float().to(weight.device), w2b.float().to(weight.device))
499

500
501
502
503
504
                try:
                    weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)

505
        return weight
506

comfyanonymous's avatar
comfyanonymous committed
507
    def unpatch_model(self, device_to=None):
508
        keys = list(self.backup.keys())
509

510
        for k in keys:
511
            set_attr(self.model, k, self.backup[k])
512

513
514
        self.backup = {}

comfyanonymous's avatar
comfyanonymous committed
515
516
517
518
519
        if device_to is not None:
            self.model.to(device_to)
            self.current_device = device_to


520
def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
521
522
    key_map = model_lora_keys_unet(model.model)
    key_map = model_lora_keys_clip(clip.cond_stage_model, key_map)
523
    loaded = load_lora(lora, key_map)
524
525
526
527
528
529
530
531
532
533
534
    new_modelpatcher = model.clone()
    k = new_modelpatcher.add_patches(loaded, strength_model)
    new_clip = clip.clone()
    k1 = new_clip.add_patches(loaded, strength_clip)
    k = set(k)
    k1 = set(k1)
    for x in loaded:
        if (x not in k) and (x not in k1):
            print("NOT LOADED", x)

    return (new_modelpatcher, new_clip)
comfyanonymous's avatar
comfyanonymous committed
535
536
537


class CLIP:
538
    def __init__(self, target=None, embedding_directory=None, no_init=False):
539
540
        if no_init:
            return
comfyanonymous's avatar
comfyanonymous committed
541
        params = target.params.copy()
542
543
        clip = target.clip
        tokenizer = target.tokenizer
544

545
546
        load_device = model_management.text_encoder_device()
        offload_device = model_management.text_encoder_offload_device()
comfyanonymous's avatar
comfyanonymous committed
547
        params['device'] = load_device
548
        if model_management.should_use_fp16(load_device, prioritize_performance=False):
549
550
551
552
553
            params['dtype'] = torch.float16
        else:
            params['dtype'] = torch.float32

        self.cond_stage_model = clip(**(params))
554

555
        self.tokenizer = tokenizer(embedding_directory=embedding_directory)
556
        self.patcher = ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
557
        self.layer_idx = None
558
559
560
561
562
563

    def clone(self):
        n = CLIP(no_init=True)
        n.patcher = self.patcher.clone()
        n.cond_stage_model = self.cond_stage_model
        n.tokenizer = self.tokenizer
comfyanonymous's avatar
comfyanonymous committed
564
        n.layer_idx = self.layer_idx
565
566
        return n

567
    def load_from_state_dict(self, sd):
568
        self.cond_stage_model.load_sd(sd)
569

570
571
    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
        return self.patcher.add_patches(patches, strength_patch, strength_model)
comfyanonymous's avatar
comfyanonymous committed
572

573
    def clip_layer(self, layer_idx):
comfyanonymous's avatar
comfyanonymous committed
574
        self.layer_idx = layer_idx
575

576
577
    def tokenize(self, text, return_word_ids=False):
        return self.tokenizer.tokenize_with_weights(text, return_word_ids)
BlenderNeko's avatar
BlenderNeko committed
578

579
    def encode_from_tokens(self, tokens, return_pooled=False):
580
581
        if self.layer_idx is not None:
            self.cond_stage_model.clip_layer(self.layer_idx)
582
583
        else:
            self.cond_stage_model.reset_clip_layer()
584

585
        self.load_model()
586
        cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
587
        if return_pooled:
588
589
            return cond, pooled
        return cond
comfyanonymous's avatar
comfyanonymous committed
590

591
    def encode(self, text):
592
        tokens = self.tokenize(text)
593
594
        return self.encode_from_tokens(tokens)

595
596
    def load_sd(self, sd):
        return self.cond_stage_model.load_sd(sd)
597

598
599
600
    def get_sd(self):
        return self.cond_stage_model.state_dict()

601
602
603
    def load_model(self):
        model_management.load_model_gpu(self.patcher)
        return self.patcher
604

605
606
607
    def get_key_patches(self):
        return self.patcher.get_key_patches()

comfyanonymous's avatar
comfyanonymous committed
608
class VAE:
609
    def __init__(self, ckpt_path=None, device=None, config=None):
comfyanonymous's avatar
comfyanonymous committed
610
611
612
        if config is None:
            #default SD1.x/SD2.x VAE parameters
            ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
613
            self.first_stage_model = AutoencoderKL(ddconfig, {'target': 'torch.nn.Identity'}, 4, monitor="val/rec_loss")
comfyanonymous's avatar
comfyanonymous committed
614
        else:
615
            self.first_stage_model = AutoencoderKL(**(config['params']))
comfyanonymous's avatar
comfyanonymous committed
616
        self.first_stage_model = self.first_stage_model.eval()
617
618
619
620
621
622
        if ckpt_path is not None:
            sd = utils.load_torch_file(ckpt_path)
            if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
                sd = diffusers_convert.convert_vae_state_dict(sd)
            self.first_stage_model.load_state_dict(sd, strict=False)

623
        if device is None:
624
            device = model_management.vae_device()
comfyanonymous's avatar
comfyanonymous committed
625
        self.device = device
626
        self.offload_device = model_management.vae_offload_device()
627
628
        self.vae_dtype = model_management.vae_dtype()
        self.first_stage_model.to(self.vae_dtype)
comfyanonymous's avatar
comfyanonymous committed
629

630
    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
pythongosssss's avatar
pythongosssss committed
631
        steps = samples.shape[0] * utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
comfyanonymous's avatar
comfyanonymous committed
632
633
        steps += samples.shape[0] * utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
        steps += samples.shape[0] * utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
pythongosssss's avatar
pythongosssss committed
634
        pbar = utils.ProgressBar(steps)
635

636
        decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
637
        output = torch.clamp((
638
639
640
            (utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = 8, pbar = pbar) +
            utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = 8, pbar = pbar) +
             utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = 8, pbar = pbar))
641
642
643
            / 3.0) / 2.0, min=0.0, max=1.0)
        return output

644
645
646
647
648
649
    def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
        steps = pixel_samples.shape[0] * utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
        steps += pixel_samples.shape[0] * utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
        steps += pixel_samples.shape[0] * utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)
        pbar = utils.ProgressBar(steps)

650
        encode_fn = lambda a: self.first_stage_model.encode(2. * a.to(self.vae_dtype).to(self.device) - 1.).sample().float()
651
652
653
654
655
656
        samples = utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
        samples += utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
        samples += utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
        samples /= 3.0
        return samples

657
    def decode(self, samples_in):
comfyanonymous's avatar
comfyanonymous committed
658
        self.first_stage_model = self.first_stage_model.to(self.device)
659
        try:
660
            memory_used = (2562 * samples_in.shape[2] * samples_in.shape[3] * 64) * 1.7
comfyanonymous's avatar
comfyanonymous committed
661
            model_management.free_memory(memory_used, self.device)
662
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
663
            batch_number = int(free_memory / memory_used)
664
665
666
667
            batch_number = max(1, batch_number)

            pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * 8), round(samples_in.shape[3] * 8)), device="cpu")
            for x in range(0, samples_in.shape[0], batch_number):
668
669
                samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
                pixel_samples[x:x+batch_number] = torch.clamp((self.first_stage_model.decode(samples) + 1.0) / 2.0, min=0.0, max=1.0).cpu().float()
670
671
672
673
        except model_management.OOM_EXCEPTION as e:
            print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
            pixel_samples = self.decode_tiled_(samples_in)

674
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
675
676
677
        pixel_samples = pixel_samples.cpu().movedim(1,-1)
        return pixel_samples

678
    def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
679
        self.first_stage_model = self.first_stage_model.to(self.device)
680
        output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
681
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
682
683
        return output.movedim(1,-1)

comfyanonymous's avatar
comfyanonymous committed
684
685
    def encode(self, pixel_samples):
        self.first_stage_model = self.first_stage_model.to(self.device)
686
687
        pixel_samples = pixel_samples.movedim(-1,1)
        try:
688
            memory_used = (2078 * pixel_samples.shape[2] * pixel_samples.shape[3]) * 1.7 #NOTE: this constant along with the one in the decode above are estimated from the mem usage for the VAE and could change.
comfyanonymous's avatar
comfyanonymous committed
689
            model_management.free_memory(memory_used, self.device)
690
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
691
            batch_number = int(free_memory / memory_used)
692
            batch_number = max(1, batch_number)
693
694
            samples = torch.empty((pixel_samples.shape[0], 4, round(pixel_samples.shape[2] // 8), round(pixel_samples.shape[3] // 8)), device="cpu")
            for x in range(0, pixel_samples.shape[0], batch_number):
695
696
                pixels_in = (2. * pixel_samples[x:x+batch_number] - 1.).to(self.vae_dtype).to(self.device)
                samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).sample().cpu().float()
697

698
699
700
701
        except model_management.OOM_EXCEPTION as e:
            print("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
            samples = self.encode_tiled_(pixel_samples)

702
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
703
704
        return samples

comfyanonymous's avatar
comfyanonymous committed
705
706
    def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
        self.first_stage_model = self.first_stage_model.to(self.device)
707
708
        pixel_samples = pixel_samples.movedim(-1,1)
        samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
709
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
710
        return samples
711

712
713
714
715
    def get_sd(self):
        return self.first_stage_model.state_dict()


BlenderNeko's avatar
BlenderNeko committed
716
def broadcast_image_to(tensor, target_batch_size, batched_number):
717
    current_batch_size = tensor.shape[0]
718
    #print(current_batch_size, target_batch_size)
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
    if current_batch_size == 1:
        return tensor

    per_batch = target_batch_size // batched_number
    tensor = tensor[:per_batch]

    if per_batch > tensor.shape[0]:
        tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0)

    current_batch_size = tensor.shape[0]
    if current_batch_size == target_batch_size:
        return tensor
    else:
        return torch.cat([tensor] * batched_number, dim=0)

734
735
class ControlBase:
    def __init__(self, device=None):
comfyanonymous's avatar
comfyanonymous committed
736
737
        self.cond_hint_original = None
        self.cond_hint = None
738
        self.strength = 1.0
739
740
741
        self.timestep_percent_range = (1.0, 0.0)
        self.timestep_range = None

742
743
        if device is None:
            device = model_management.get_torch_device()
744
        self.device = device
comfyanonymous's avatar
comfyanonymous committed
745
        self.previous_controlnet = None
746
        self.global_average_pooling = False
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781

    def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(1.0, 0.0)):
        self.cond_hint_original = cond_hint
        self.strength = strength
        self.timestep_percent_range = timestep_percent_range
        return self

    def pre_run(self, model, percent_to_timestep_function):
        self.timestep_range = (percent_to_timestep_function(self.timestep_percent_range[0]), percent_to_timestep_function(self.timestep_percent_range[1]))
        if self.previous_controlnet is not None:
            self.previous_controlnet.pre_run(model, percent_to_timestep_function)

    def set_previous_controlnet(self, controlnet):
        self.previous_controlnet = controlnet
        return self

    def cleanup(self):
        if self.previous_controlnet is not None:
            self.previous_controlnet.cleanup()
        if self.cond_hint is not None:
            del self.cond_hint
            self.cond_hint = None
        self.timestep_range = None

    def get_models(self):
        out = []
        if self.previous_controlnet is not None:
            out += self.previous_controlnet.get_models()
        return out

    def copy_to(self, c):
        c.cond_hint_original = self.cond_hint_original
        c.strength = self.strength
        c.timestep_percent_range = self.timestep_percent_range

782
783
784
785
786
    def inference_memory_requirements(self, dtype):
        if self.previous_controlnet is not None:
            return self.previous_controlnet.inference_memory_requirements(dtype)
        return 0

787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
    def control_merge(self, control_input, control_output, control_prev, output_dtype):
        out = {'input':[], 'middle':[], 'output': []}

        if control_input is not None:
            for i in range(len(control_input)):
                key = 'input'
                x = control_input[i]
                if x is not None:
                    x *= self.strength
                    if x.dtype != output_dtype:
                        x = x.to(output_dtype)
                out[key].insert(0, x)

        if control_output is not None:
            for i in range(len(control_output)):
                if i == (len(control_output) - 1):
                    key = 'middle'
                    index = 0
                else:
                    key = 'output'
                    index = i
                x = control_output[i]
                if x is not None:
                    if self.global_average_pooling:
                        x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])

                    x *= self.strength
                    if x.dtype != output_dtype:
                        x = x.to(output_dtype)

                out[key].append(x)
        if control_prev is not None:
            for x in ['input', 'middle', 'output']:
                o = out[x]
                for i in range(len(control_prev[x])):
                    prev_val = control_prev[x][i]
                    if i >= len(o):
                        o.append(prev_val)
                    elif prev_val is not None:
                        if o[i] is None:
                            o[i] = prev_val
                        else:
                            o[i] += prev_val
        return out

832
833
834
835
class ControlNet(ControlBase):
    def __init__(self, control_model, global_average_pooling=False, device=None):
        super().__init__(device)
        self.control_model = control_model
comfyanonymous's avatar
comfyanonymous committed
836
        self.control_model_wrapped = ModelPatcher(self.control_model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
837
        self.global_average_pooling = global_average_pooling
comfyanonymous's avatar
comfyanonymous committed
838

839
    def get_control(self, x_noisy, t, cond, batched_number):
comfyanonymous's avatar
comfyanonymous committed
840
841
        control_prev = None
        if self.previous_controlnet is not None:
842
            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
comfyanonymous's avatar
comfyanonymous committed
843

844
845
846
847
848
849
850
        if self.timestep_range is not None:
            if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
                if control_prev is not None:
                    return control_prev
                else:
                    return {}

851
        output_dtype = x_noisy.dtype
comfyanonymous's avatar
comfyanonymous committed
852
853
854
855
        if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
            if self.cond_hint is not None:
                del self.cond_hint
            self.cond_hint = None
BlenderNeko's avatar
BlenderNeko committed
856
857
858
            self.cond_hint = utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(self.control_model.dtype).to(self.device)
        if x_noisy.shape[0] != self.cond_hint.shape[0]:
            self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
859
860


861
862
863
864
865
        context = torch.cat(cond['c_crossattn'], 1)
        y = cond.get('c_adm', None)
        if y is not None:
            y = y.to(self.control_model.dtype)
        control = self.control_model(x=x_noisy.to(self.control_model.dtype), hint=self.cond_hint, timesteps=t, context=context.to(self.control_model.dtype), y=y)
866
        return self.control_merge(None, control, control_prev, output_dtype)
comfyanonymous's avatar
comfyanonymous committed
867
868

    def copy(self):
869
        c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling)
870
        self.copy_to(c)
comfyanonymous's avatar
comfyanonymous committed
871
872
        return c

873
874
    def get_models(self):
        out = super().get_models()
comfyanonymous's avatar
comfyanonymous committed
875
        out.append(self.control_model_wrapped)
876
877
        return out

comfyanonymous's avatar
comfyanonymous committed
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
class ControlLoraOps:
    class Linear(torch.nn.Module):
        def __init__(self, in_features: int, out_features: int, bias: bool = True,
                    device=None, dtype=None) -> None:
            factory_kwargs = {'device': device, 'dtype': dtype}
            super().__init__()
            self.in_features = in_features
            self.out_features = out_features
            self.weight = None
            self.up = None
            self.down = None
            self.bias = None

        def forward(self, input):
            if self.up is not None:
comfyanonymous's avatar
comfyanonymous committed
893
                return torch.nn.functional.linear(input, self.weight.to(input.device) + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), self.bias)
comfyanonymous's avatar
comfyanonymous committed
894
            else:
comfyanonymous's avatar
comfyanonymous committed
895
                return torch.nn.functional.linear(input, self.weight.to(input.device), self.bias)
comfyanonymous's avatar
comfyanonymous committed
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931

    class Conv2d(torch.nn.Module):
        def __init__(
            self,
            in_channels,
            out_channels,
            kernel_size,
            stride=1,
            padding=0,
            dilation=1,
            groups=1,
            bias=True,
            padding_mode='zeros',
            device=None,
            dtype=None
        ):
            super().__init__()
            self.in_channels = in_channels
            self.out_channels = out_channels
            self.kernel_size = kernel_size
            self.stride = stride
            self.padding = padding
            self.dilation = dilation
            self.transposed = False
            self.output_padding = 0
            self.groups = groups
            self.padding_mode = padding_mode

            self.weight = None
            self.bias = None
            self.up = None
            self.down = None


        def forward(self, input):
            if self.up is not None:
comfyanonymous's avatar
comfyanonymous committed
932
                return torch.nn.functional.conv2d(input, self.weight.to(input.device) + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), self.bias, self.stride, self.padding, self.dilation, self.groups)
comfyanonymous's avatar
comfyanonymous committed
933
            else:
comfyanonymous's avatar
comfyanonymous committed
934
                return torch.nn.functional.conv2d(input, self.weight.to(input.device), self.bias, self.stride, self.padding, self.dilation, self.groups)
comfyanonymous's avatar
comfyanonymous committed
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955

    def conv_nd(self, dims, *args, **kwargs):
        if dims == 2:
            return self.Conv2d(*args, **kwargs)
        else:
            raise ValueError(f"unsupported dimensions: {dims}")


class ControlLora(ControlNet):
    def __init__(self, control_weights, global_average_pooling=False, device=None):
        ControlBase.__init__(self, device)
        self.control_weights = control_weights
        self.global_average_pooling = global_average_pooling

    def pre_run(self, model, percent_to_timestep_function):
        super().pre_run(model, percent_to_timestep_function)
        controlnet_config = model.model_config.unet_config.copy()
        controlnet_config.pop("out_channels")
        controlnet_config["hint_channels"] = self.control_weights["input_hint_block.0.weight"].shape[1]
        controlnet_config["operations"] = ControlLoraOps()
        self.control_model = cldm.ControlNet(**controlnet_config)
956
957
        dtype = model.get_dtype()
        self.control_model.to(dtype)
comfyanonymous's avatar
comfyanonymous committed
958
959
960
961
962
963
        self.control_model.to(model_management.get_torch_device())
        diffusion_model = model.diffusion_model
        sd = diffusion_model.state_dict()
        cm = self.control_model.state_dict()

        for k in sd:
comfyanonymous's avatar
comfyanonymous committed
964
965
966
967
968
969
            weight = sd[k]
            if weight.device == torch.device("meta"): #lowvram NOTE: this depends on the inner working of the accelerate library so it might break.
                key_split = k.split('.')              # I have no idea why they don't just leave the weight there instead of using the meta device.
                op = get_attr(diffusion_model, '.'.join(key_split[:-1]))
                weight = op._hf_hook.weights_map[key_split[-1]]

comfyanonymous's avatar
comfyanonymous committed
970
            try:
comfyanonymous's avatar
comfyanonymous committed
971
                set_attr(self.control_model, k, weight)
comfyanonymous's avatar
comfyanonymous committed
972
973
974
975
976
            except:
                pass

        for k in self.control_weights:
            if k not in {"lora_controlnet"}:
977
                set_attr(self.control_model, k, self.control_weights[k].to(dtype).to(model_management.get_torch_device()))
comfyanonymous's avatar
comfyanonymous committed
978
979
980
981
982
983
984
985
986
987
988
989
990
991

    def copy(self):
        c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)
        self.copy_to(c)
        return c

    def cleanup(self):
        del self.control_model
        self.control_model = None
        super().cleanup()

    def get_models(self):
        out = ControlBase.get_models(self)
        return out
992

993
994
995
    def inference_memory_requirements(self, dtype):
        return utils.calculate_parameters(self.control_weights) * model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)

996
def load_controlnet(ckpt_path, model=None):
997
    controlnet_data = utils.load_torch_file(ckpt_path, safe_load=True)
comfyanonymous's avatar
comfyanonymous committed
998
999
    if "lora_controlnet" in controlnet_data:
        return ControlLora(controlnet_data)
1000
1001
1002
1003

    controlnet_config = None
    if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
        use_fp16 = model_management.should_use_fp16()
1004
        controlnet_config = model_detection.unet_config_from_diffusers_unet(controlnet_data, use_fp16)
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
        diffusers_keys = utils.unet_to_diffusers(controlnet_config)
        diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
        diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"

        count = 0
        loop = True
        while loop:
            suffix = [".weight", ".bias"]
            for s in suffix:
                k_in = "controlnet_down_blocks.{}{}".format(count, s)
                k_out = "zero_convs.{}.0{}".format(count, s)
                if k_in not in controlnet_data:
                    loop = False
                    break
                diffusers_keys[k_in] = k_out
            count += 1

        count = 0
        loop = True
        while loop:
            suffix = [".weight", ".bias"]
            for s in suffix:
                if count == 0:
                    k_in = "controlnet_cond_embedding.conv_in{}".format(s)
                else:
                    k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
                k_out = "input_hint_block.{}{}".format(count * 2, s)
                if k_in not in controlnet_data:
                    k_in = "controlnet_cond_embedding.conv_out{}".format(s)
                    loop = False
                diffusers_keys[k_in] = k_out
            count += 1

        new_sd = {}
        for k in diffusers_keys:
            if k in controlnet_data:
                new_sd[diffusers_keys[k]] = controlnet_data.pop(k)

1043
1044
1045
        leftover_keys = controlnet_data.keys()
        if len(leftover_keys) > 0:
            print("leftover keys:", leftover_keys)
1046
1047
        controlnet_data = new_sd

1048
    pth_key = 'control_model.zero_convs.0.0.weight'
comfyanonymous's avatar
comfyanonymous committed
1049
    pth = False
1050
    key = 'zero_convs.0.0.weight'
comfyanonymous's avatar
comfyanonymous committed
1051
1052
1053
    if pth_key in controlnet_data:
        pth = True
        key = pth_key
1054
        prefix = "control_model."
comfyanonymous's avatar
comfyanonymous committed
1055
    elif key in controlnet_data:
1056
        prefix = ""
comfyanonymous's avatar
comfyanonymous committed
1057
    else:
1058
1059
1060
1061
        net = load_t2i_adapter(controlnet_data)
        if net is None:
            print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path)
        return net
comfyanonymous's avatar
comfyanonymous committed
1062

1063
1064
1065
    if controlnet_config is None:
        use_fp16 = model_management.should_use_fp16()
        controlnet_config = model_detection.model_config_from_unet(controlnet_data, prefix, use_fp16).unet_config
1066
    controlnet_config.pop("out_channels")
1067
    controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
1068
1069
    control_model = cldm.ControlNet(**controlnet_config)

comfyanonymous's avatar
comfyanonymous committed
1070
    if pth:
1071
1072
        if 'difference' in controlnet_data:
            if model is not None:
1073
1074
                model_management.load_models_gpu([model])
                model_sd = model.model_state_dict()
1075
1076
1077
                for x in controlnet_data:
                    c_m = "control_model."
                    if x.startswith(c_m):
comfyanonymous's avatar
comfyanonymous committed
1078
                        sd_key = "diffusion_model.{}".format(x[len(c_m):])
1079
1080
1081
1082
1083
1084
                        if sd_key in model_sd:
                            cd = controlnet_data[x]
                            cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
            else:
                print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")

comfyanonymous's avatar
comfyanonymous committed
1085
1086
1087
1088
        class WeightsLoader(torch.nn.Module):
            pass
        w = WeightsLoader()
        w.control_model = control_model
1089
        missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
comfyanonymous's avatar
comfyanonymous committed
1090
    else:
1091
1092
        missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
    print(missing, unexpected)
comfyanonymous's avatar
comfyanonymous committed
1093

1094
1095
1096
    if use_fp16:
        control_model = control_model.half()

1097
1098
1099
1100
1101
    global_average_pooling = False
    if ckpt_path.endswith("_shuffle.pth") or ckpt_path.endswith("_shuffle.safetensors") or ckpt_path.endswith("_shuffle_fp16.safetensors"): #TODO: smarter way of enabling global_average_pooling
        global_average_pooling = True

    control = ControlNet(control_model, global_average_pooling=global_average_pooling)
comfyanonymous's avatar
comfyanonymous committed
1102
1103
    return control

1104
class T2IAdapter(ControlBase):
1105
    def __init__(self, t2i_model, channels_in, device=None):
1106
        super().__init__(device)
1107
1108
1109
1110
        self.t2i_model = t2i_model
        self.channels_in = channels_in
        self.control_input = None

1111
1112
1113
1114
1115
1116
    def scale_image_to(self, width, height):
        unshuffle_amount = self.t2i_model.unshuffle_amount
        width = math.ceil(width / unshuffle_amount) * unshuffle_amount
        height = math.ceil(height / unshuffle_amount) * unshuffle_amount
        return width, height

1117
    def get_control(self, x_noisy, t, cond, batched_number):
1118
1119
        control_prev = None
        if self.previous_controlnet is not None:
1120
            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
1121

1122
1123
1124
1125
1126
1127
1128
        if self.timestep_range is not None:
            if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
                if control_prev is not None:
                    return control_prev
                else:
                    return {}

1129
1130
1131
        if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
            if self.cond_hint is not None:
                del self.cond_hint
BlenderNeko's avatar
BlenderNeko committed
1132
            self.control_input = None
1133
            self.cond_hint = None
1134
1135
            width, height = self.scale_image_to(x_noisy.shape[3] * 8, x_noisy.shape[2] * 8)
            self.cond_hint = utils.common_upscale(self.cond_hint_original, width, height, 'nearest-exact', "center").float().to(self.device)
1136
1137
            if self.channels_in == 1 and self.cond_hint.shape[1] > 1:
                self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)
BlenderNeko's avatar
BlenderNeko committed
1138
1139
1140
        if x_noisy.shape[0] != self.cond_hint.shape[0]:
            self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
        if self.control_input is None:
1141
            self.t2i_model.to(x_noisy.dtype)
1142
            self.t2i_model.to(self.device)
1143
            self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
1144
1145
            self.t2i_model.cpu()

1146
        control_input = list(map(lambda a: None if a is None else a.clone(), self.control_input))
comfyanonymous's avatar
comfyanonymous committed
1147
1148
1149
1150
1151
        mid = None
        if self.t2i_model.xl == True:
            mid = control_input[-1:]
            control_input = control_input[:-1]
        return self.control_merge(control_input, mid, control_prev, x_noisy.dtype)
1152
1153
1154

    def copy(self):
        c = T2IAdapter(self.t2i_model, self.channels_in)
1155
        self.copy_to(c)
1156
1157
        return c

1158
def load_t2i_adapter(t2i_data):
1159
    keys = t2i_data.keys()
1160
1161
1162
    if 'adapter' in keys:
        t2i_data = t2i_data['adapter']
        keys = t2i_data.keys()
1163
    if "body.0.in_conv.weight" in keys:
1164
1165
        cin = t2i_data['body.0.in_conv.weight'].shape[1]
        model_ad = adapter.Adapter_light(cin=cin, channels=[320, 640, 1280, 1280], nums_rb=4)
1166
    elif 'conv_in.weight' in keys:
1167
        cin = t2i_data['conv_in.weight'].shape[1]
1168
1169
1170
1171
1172
1173
        channel = t2i_data['conv_in.weight'].shape[0]
        ksize = t2i_data['body.0.block2.weight'].shape[2]
        use_conv = False
        down_opts = list(filter(lambda a: a.endswith("down_opt.op.weight"), keys))
        if len(down_opts) > 0:
            use_conv = True
comfyanonymous's avatar
comfyanonymous committed
1174
1175
1176
1177
        xl = False
        if cin == 256:
            xl = True
        model_ad = adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)
1178
1179
    else:
        return None
comfyanonymous's avatar
comfyanonymous committed
1180
1181
1182
1183
1184
1185
1186
1187
    missing, unexpected = model_ad.load_state_dict(t2i_data)
    if len(missing) > 0:
        print("t2i missing", missing)

    if len(unexpected) > 0:
        print("t2i unexpected", unexpected)

    return T2IAdapter(model_ad, model_ad.input_channels)
comfyanonymous's avatar
comfyanonymous committed
1188

1189
1190
1191
1192
1193
1194
1195
1196
1197
1198

class StyleModel:
    def __init__(self, model, device="cpu"):
        self.model = model

    def get_cond(self, input):
        return self.model(input.last_hidden_state)


def load_style_model(ckpt_path):
1199
    model_data = utils.load_torch_file(ckpt_path, safe_load=True)
1200
1201
1202
1203
1204
1205
1206
1207
1208
    keys = model_data.keys()
    if "style_embedding" in keys:
        model = adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
    else:
        raise Exception("invalid style model {}".format(ckpt_path))
    model.load_state_dict(model_data)
    return StyleModel(model)


1209
1210
1211
1212
1213
def load_clip(ckpt_paths, embedding_directory=None):
    clip_data = []
    for p in ckpt_paths:
        clip_data.append(utils.load_torch_file(p, safe_load=True))

comfyanonymous's avatar
comfyanonymous committed
1214
1215
1216
    class EmptyClass:
        pass

1217
1218
1219
1220
    for i in range(len(clip_data)):
        if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
            clip_data[i] = utils.transformers_convert(clip_data[i], "", "text_model.", 32)

comfyanonymous's avatar
comfyanonymous committed
1221
1222
    clip_target = EmptyClass()
    clip_target.params = {}
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
    if len(clip_data) == 1:
        if "text_model.encoder.layers.30.mlp.fc1.weight" in clip_data[0]:
            clip_target.clip = sdxl_clip.SDXLRefinerClipModel
            clip_target.tokenizer = sdxl_clip.SDXLTokenizer
        elif "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data[0]:
            clip_target.clip = sd2_clip.SD2ClipModel
            clip_target.tokenizer = sd2_clip.SD2Tokenizer
        else:
            clip_target.clip = sd1_clip.SD1ClipModel
            clip_target.tokenizer = sd1_clip.SD1Tokenizer
1233
    else:
1234
1235
        clip_target.clip = sdxl_clip.SDXLClipModel
        clip_target.tokenizer = sdxl_clip.SDXLTokenizer
comfyanonymous's avatar
comfyanonymous committed
1236
1237

    clip = CLIP(clip_target, embedding_directory=embedding_directory)
1238
1239
1240
1241
1242
1243
1244
    for c in clip_data:
        m, u = clip.load_sd(c)
        if len(m) > 0:
            print("clip missing:", m)

        if len(u) > 0:
            print("clip unexpected:", u)
1245
    return clip
comfyanonymous's avatar
comfyanonymous committed
1246

1247
def load_gligen(ckpt_path):
1248
    data = utils.load_torch_file(ckpt_path, safe_load=True)
1249
1250
1251
    model = gligen.load_gligen(data)
    if model_management.should_use_fp16():
        model = model.half()
comfyanonymous's avatar
comfyanonymous committed
1252
    return ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
1253

comfyanonymous's avatar
comfyanonymous committed
1254
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
1255
    #TODO: this function is a mess and should be removed eventually
comfyanonymous's avatar
comfyanonymous committed
1256
1257
1258
    if config is None:
        with open(config_path, 'r') as stream:
            config = yaml.safe_load(stream)
comfyanonymous's avatar
comfyanonymous committed
1259
1260
1261
1262
1263
    model_config_params = config['model']['params']
    clip_config = model_config_params['cond_stage_config']
    scale_factor = model_config_params['scale_factor']
    vae_config = model_config_params['first_stage_config']

1264
1265
1266
    fp16 = False
    if "unet_config" in model_config_params:
        if "params" in model_config_params["unet_config"]:
comfyanonymous's avatar
comfyanonymous committed
1267
1268
1269
1270
1271
1272
1273
1274
            unet_config = model_config_params["unet_config"]["params"]
            if "use_fp16" in unet_config:
                fp16 = unet_config["use_fp16"]

    noise_aug_config = None
    if "noise_aug_config" in model_config_params:
        noise_aug_config = model_config_params["noise_aug_config"]

1275
    model_type = model_base.ModelType.EPS
comfyanonymous's avatar
comfyanonymous committed
1276
1277
1278

    if "parameterization" in model_config_params:
        if model_config_params["parameterization"] == "v":
1279
            model_type = model_base.ModelType.V_PREDICTION
1280

comfyanonymous's avatar
comfyanonymous committed
1281
1282
1283
1284
1285
1286
    clip = None
    vae = None

    class WeightsLoader(torch.nn.Module):
        pass

1287
1288
    if state_dict is None:
        state_dict = utils.load_torch_file(ckpt_path)
comfyanonymous's avatar
comfyanonymous committed
1289

1290
1291
1292
1293
1294
1295
1296
1297
    class EmptyClass:
        pass

    model_config = EmptyClass()
    model_config.unet_config = unet_config
    from . import latent_formats
    model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)

comfyanonymous's avatar
comfyanonymous committed
1298
    if config['model']["target"].endswith("LatentInpaintDiffusion"):
1299
        model = model_base.SDInpaint(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1300
    elif config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
1301
        model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1302
    else:
1303
        model = model_base.BaseModel(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1304

1305
1306
1307
    if fp16:
        model = model.half()

1308
1309
    offload_device = model_management.unet_offload_device()
    model = model.to(offload_device)
1310
1311
1312
1313
    model.load_model_weights(state_dict, "model.diffusion_model.")

    if output_vae:
        w = WeightsLoader()
1314
        vae = VAE(config=vae_config)
1315
1316
1317
1318
1319
1320
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, state_dict)

    if output_clip:
        w = WeightsLoader()
        clip_target = EmptyClass()
1321
        clip_target.params = clip_config.get("params", {})
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
        if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
            clip_target.clip = sd2_clip.SD2ClipModel
            clip_target.tokenizer = sd2_clip.SD2Tokenizer
        elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
            clip_target.clip = sd1_clip.SD1ClipModel
            clip_target.tokenizer = sd1_clip.SD1Tokenizer
        clip = CLIP(clip_target, embedding_directory=embedding_directory)
        w.cond_stage_model = clip.cond_stage_model
        load_clip_weights(w, state_dict)

1332
    return (ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
1333

1334
1335
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None):
    sd = utils.load_torch_file(ckpt_path)
1336
1337
    sd_keys = sd.keys()
    clip = None
1338
    clipvision = None
1339
    vae = None
1340
1341
    model = None
    clip_target = None
1342

1343
    parameters = utils.calculate_parameters(sd, "model.diffusion_model.")
1344
    fp16 = model_management.should_use_fp16(model_params=parameters)
1345

1346
1347
1348
    class WeightsLoader(torch.nn.Module):
        pass

1349
1350
1351
    model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", fp16)
    if model_config is None:
        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
1352

1353
    if model_config.clip_vision_prefix is not None:
1354
        if output_clipvision:
1355
            clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
1356

comfyanonymous's avatar
comfyanonymous committed
1357
1358
1359
1360
1361
    dtype = torch.float32
    if fp16:
        dtype = torch.float16

    inital_load_device = model_management.unet_inital_load_device(parameters, dtype)
1362
    offload_device = model_management.unet_offload_device()
comfyanonymous's avatar
comfyanonymous committed
1363
    model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
1364
    model.load_model_weights(sd, "model.diffusion_model.")
1365

1366
    if output_vae:
1367
        vae = VAE()
1368
1369
1370
        w = WeightsLoader()
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, sd)
1371

1372
1373
1374
1375
1376
1377
1378
    if output_clip:
        w = WeightsLoader()
        clip_target = model_config.clip_target()
        clip = CLIP(clip_target, embedding_directory=embedding_directory)
        w.cond_stage_model = clip.cond_stage_model
        sd = model_config.process_clip_state_dict(sd)
        load_model_weights(w, sd)
comfyanonymous's avatar
comfyanonymous committed
1379

1380
1381
1382
    left_over = sd.keys()
    if len(left_over) > 0:
        print("left over keys:", left_over)
1383

comfyanonymous's avatar
comfyanonymous committed
1384
1385
1386
1387
1388
1389
    model_patcher = ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
    if inital_load_device != torch.device("cpu"):
        print("loaded straight to GPU")
        model_management.load_model_gpu(model_patcher)

    return (model_patcher, clip, vae, clipvision)
1390

1391
1392
1393

def load_unet(unet_path): #load unet in diffusers format
    sd = utils.load_torch_file(unet_path)
1394
    parameters = utils.calculate_parameters(sd)
1395
1396
    fp16 = model_management.should_use_fp16(model_params=parameters)

1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
    model_config = model_detection.model_config_from_diffusers_unet(sd, fp16)
    if model_config is None:
        print("ERROR UNSUPPORTED UNET", unet_path)
        return None

    diffusers_keys = utils.unet_to_diffusers(model_config.unet_config)

    new_sd = {}
    for k in diffusers_keys:
        if k in sd:
            new_sd[diffusers_keys[k]] = sd.pop(k)
        else:
            print(diffusers_keys[k], k)
    offload_device = model_management.unet_offload_device()
    model = model_config.get_model(new_sd, "")
    model = model.to(offload_device)
    model.load_model_weights(new_sd, "")
    return ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device)
1415

1416
def save_checkpoint(output_path, model, clip, vae, metadata=None):
1417
1418
1419
    model_management.load_models_gpu([model, clip.load_model()])
    sd = model.model.state_dict_for_saving(clip.get_sd(), vae.get_sd())
    utils.save_torch_file(sd, output_path, metadata=metadata)