sd.py 49.6 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
26
import comfy.lora

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

    k = list(sd.keys())
    for x in k:
34
35
36
37
38
39
40
41
42
43
        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
44
45
46
47
        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
48
49
50
51
    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()
52

53
54
    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
55

56

57

58
59
60
61
62
63
64
65
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
66
67
68
69
70
71
72
def get_attr(obj, attr):
    attrs = attr.split(".")
    for name in attrs:
        obj = getattr(obj, name)
    return obj


73
class ModelPatcher:
comfyanonymous's avatar
comfyanonymous committed
74
    def __init__(self, model, load_device, offload_device, size=0, current_device=None):
75
        self.size = size
76
        self.model = model
77
        self.patches = {}
78
        self.backup = {}
79
        self.model_options = {"transformer_options":{}}
80
        self.model_size()
81
82
        self.load_device = load_device
        self.offload_device = offload_device
comfyanonymous's avatar
comfyanonymous committed
83
84
85
86
        if current_device is None:
            self.current_device = self.offload_device
        else:
            self.current_device = current_device
87
88
89
90
91
92
93
94
95
96

    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
97
        self.model_keys = set(model_sd.keys())
98
        return size
99
100

    def clone(self):
comfyanonymous's avatar
comfyanonymous committed
101
        n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device)
102
103
104
105
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]

106
        n.model_options = copy.deepcopy(self.model_options)
107
        n.model_keys = self.model_keys
108
109
        return n

comfyanonymous's avatar
comfyanonymous committed
110
111
112
113
114
    def is_clone(self, other):
        if hasattr(other, 'model') and self.model is other.model:
            return True
        return False

115
    def set_model_sampler_cfg_function(self, sampler_cfg_function):
116
117
118
119
        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
120

121
122
123
    def set_model_unet_function_wrapper(self, unet_wrapper_function):
        self.model_options["model_function_wrapper"] = unet_wrapper_function

124
125
126
127
128
129
    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]

130
131
132
133
134
135
136
137
    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

138
139
140
141
142
143
    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")

144
145
146
147
148
149
150
151
152
    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")

153
154
155
    def set_model_attn2_output_patch(self, patch):
        self.set_model_patch(patch, "attn2_output_patch")

156
157
158
159
160
161
162
163
164
    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)
165
166
167
168
169
170
171
        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)
172

173
    def model_dtype(self):
174
175
        if hasattr(self.model, "get_dtype"):
            return self.model.get_dtype()
176

177
    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
178
        p = set()
179
        for k in patches:
180
            if k in self.model_keys:
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
                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
200

201
    def model_state_dict(self, filter_prefix=None):
202
203
        sd = self.model.state_dict()
        keys = list(sd.keys())
204
205
206
207
        if filter_prefix is not None:
            for k in keys:
                if not k.startswith(filter_prefix):
                    sd.pop(k)
208
209
        return sd

210
    def patch_model(self, device_to=None):
211
        model_sd = self.model_state_dict()
212
213
214
215
        for key in self.patches:
            if key not in model_sd:
                print("could not patch. key doesn't exist in model:", k)
                continue
216

217
            weight = model_sd[key]
218

219
            if key not in self.backup:
220
                self.backup[key] = weight.to(self.offload_device)
221

222
223
224
225
            if device_to is not None:
                temp_weight = weight.float().to(device_to, copy=True)
            else:
                temp_weight = weight.to(torch.float32, copy=True)
226
227
            out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
            set_attr(self.model, key, out_weight)
228
            del temp_weight
comfyanonymous's avatar
comfyanonymous committed
229
230
231
232
233

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

234
        return self.model
comfyanonymous's avatar
comfyanonymous committed
235

236
237
238
239
240
241
242
243
244
245
246
247
248
249
    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]
250
251
252
253
254
                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)
255
            elif len(v) == 4: #lora/locon
comfyanonymous's avatar
comfyanonymous committed
256
257
                mat1 = v[0].float().to(weight.device)
                mat2 = v[1].float().to(weight.device)
258
259
260
261
                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
262
263
264
                    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)
265
266
267
268
                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)
269
270
271
272
273
274
275
276
277
278
279
280
281
            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
282
283
                else:
                    w1 = w1.float().to(weight.device)
284
285
286
287

                if w2 is None:
                    dim = w2_b.shape[0]
                    if t2 is None:
comfyanonymous's avatar
comfyanonymous committed
288
                        w2 = torch.mm(w2_a.float().to(weight.device), w2_b.float().to(weight.device))
289
                    else:
comfyanonymous's avatar
comfyanonymous committed
290
291
292
                        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)
293
294
295
296
297
298

                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

299
300
301
302
                try:
                    weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
303
304
305
306
307
308
309
310
311
312
            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
313
314
                    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))
315
                else:
comfyanonymous's avatar
comfyanonymous committed
316
317
                    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))
318

319
320
321
322
323
                try:
                    weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)

324
        return weight
325

comfyanonymous's avatar
comfyanonymous committed
326
    def unpatch_model(self, device_to=None):
327
        keys = list(self.backup.keys())
328

329
        for k in keys:
330
            set_attr(self.model, k, self.backup[k])
331

332
333
        self.backup = {}

comfyanonymous's avatar
comfyanonymous committed
334
335
336
337
338
        if device_to is not None:
            self.model.to(device_to)
            self.current_device = device_to


339
def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
340
341
342
    key_map = comfy.lora.model_lora_keys_unet(model.model)
    key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
    loaded = comfy.lora.load_lora(lora, key_map)
343
344
345
346
347
348
349
350
351
352
353
    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
354
355
356


class CLIP:
357
    def __init__(self, target=None, embedding_directory=None, no_init=False):
358
359
        if no_init:
            return
comfyanonymous's avatar
comfyanonymous committed
360
        params = target.params.copy()
361
362
        clip = target.clip
        tokenizer = target.tokenizer
363

364
365
        load_device = model_management.text_encoder_device()
        offload_device = model_management.text_encoder_offload_device()
comfyanonymous's avatar
comfyanonymous committed
366
        params['device'] = load_device
367
        if model_management.should_use_fp16(load_device, prioritize_performance=False):
368
369
370
371
372
            params['dtype'] = torch.float16
        else:
            params['dtype'] = torch.float32

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

374
        self.tokenizer = tokenizer(embedding_directory=embedding_directory)
375
        self.patcher = ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
376
        self.layer_idx = None
377
378
379
380
381
382

    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
383
        n.layer_idx = self.layer_idx
384
385
        return n

386
387
    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
388

389
    def clip_layer(self, layer_idx):
comfyanonymous's avatar
comfyanonymous committed
390
        self.layer_idx = layer_idx
391

392
393
    def tokenize(self, text, return_word_ids=False):
        return self.tokenizer.tokenize_with_weights(text, return_word_ids)
BlenderNeko's avatar
BlenderNeko committed
394

395
    def encode_from_tokens(self, tokens, return_pooled=False):
396
397
        if self.layer_idx is not None:
            self.cond_stage_model.clip_layer(self.layer_idx)
398
399
        else:
            self.cond_stage_model.reset_clip_layer()
400

401
        self.load_model()
402
        cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
403
        if return_pooled:
404
405
            return cond, pooled
        return cond
comfyanonymous's avatar
comfyanonymous committed
406

407
    def encode(self, text):
408
        tokens = self.tokenize(text)
409
410
        return self.encode_from_tokens(tokens)

411
412
    def load_sd(self, sd):
        return self.cond_stage_model.load_sd(sd)
413

414
415
416
    def get_sd(self):
        return self.cond_stage_model.state_dict()

417
418
419
    def load_model(self):
        model_management.load_model_gpu(self.patcher)
        return self.patcher
420

421
422
423
    def get_key_patches(self):
        return self.patcher.get_key_patches()

comfyanonymous's avatar
comfyanonymous committed
424
class VAE:
425
    def __init__(self, ckpt_path=None, device=None, config=None):
comfyanonymous's avatar
comfyanonymous committed
426
427
428
        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}
429
            self.first_stage_model = AutoencoderKL(ddconfig, {'target': 'torch.nn.Identity'}, 4, monitor="val/rec_loss")
comfyanonymous's avatar
comfyanonymous committed
430
        else:
431
            self.first_stage_model = AutoencoderKL(**(config['params']))
comfyanonymous's avatar
comfyanonymous committed
432
        self.first_stage_model = self.first_stage_model.eval()
433
434
435
436
437
438
        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)

439
        if device is None:
440
            device = model_management.vae_device()
comfyanonymous's avatar
comfyanonymous committed
441
        self.device = device
442
        self.offload_device = model_management.vae_offload_device()
443
444
        self.vae_dtype = model_management.vae_dtype()
        self.first_stage_model.to(self.vae_dtype)
comfyanonymous's avatar
comfyanonymous committed
445

446
    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
pythongosssss's avatar
pythongosssss committed
447
        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
448
449
        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
450
        pbar = utils.ProgressBar(steps)
451

452
        decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
453
        output = torch.clamp((
454
455
456
            (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))
457
458
459
            / 3.0) / 2.0, min=0.0, max=1.0)
        return output

460
461
462
463
464
465
    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)

466
        encode_fn = lambda a: self.first_stage_model.encode(2. * a.to(self.vae_dtype).to(self.device) - 1.).sample().float()
467
468
469
470
471
472
        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

473
    def decode(self, samples_in):
comfyanonymous's avatar
comfyanonymous committed
474
        self.first_stage_model = self.first_stage_model.to(self.device)
475
        try:
476
            memory_used = (2562 * samples_in.shape[2] * samples_in.shape[3] * 64) * 1.7
comfyanonymous's avatar
comfyanonymous committed
477
            model_management.free_memory(memory_used, self.device)
478
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
479
            batch_number = int(free_memory / memory_used)
480
481
482
483
            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):
484
485
                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()
486
487
488
489
        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)

490
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
491
492
493
        pixel_samples = pixel_samples.cpu().movedim(1,-1)
        return pixel_samples

494
    def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
495
        self.first_stage_model = self.first_stage_model.to(self.device)
496
        output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
497
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
498
499
        return output.movedim(1,-1)

comfyanonymous's avatar
comfyanonymous committed
500
501
    def encode(self, pixel_samples):
        self.first_stage_model = self.first_stage_model.to(self.device)
502
503
        pixel_samples = pixel_samples.movedim(-1,1)
        try:
504
            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
505
            model_management.free_memory(memory_used, self.device)
506
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
507
            batch_number = int(free_memory / memory_used)
508
            batch_number = max(1, batch_number)
509
510
            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):
511
512
                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()
513

514
515
516
517
        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)

518
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
519
520
        return samples

comfyanonymous's avatar
comfyanonymous committed
521
522
    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)
523
524
        pixel_samples = pixel_samples.movedim(-1,1)
        samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
525
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
526
        return samples
527

528
529
530
531
    def get_sd(self):
        return self.first_stage_model.state_dict()


BlenderNeko's avatar
BlenderNeko committed
532
def broadcast_image_to(tensor, target_batch_size, batched_number):
533
    current_batch_size = tensor.shape[0]
534
    #print(current_batch_size, target_batch_size)
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
    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)

550
551
class ControlBase:
    def __init__(self, device=None):
comfyanonymous's avatar
comfyanonymous committed
552
553
        self.cond_hint_original = None
        self.cond_hint = None
554
        self.strength = 1.0
555
556
557
        self.timestep_percent_range = (1.0, 0.0)
        self.timestep_range = None

558
559
        if device is None:
            device = model_management.get_torch_device()
560
        self.device = device
comfyanonymous's avatar
comfyanonymous committed
561
        self.previous_controlnet = None
562
        self.global_average_pooling = False
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

    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

598
599
600
601
602
    def inference_memory_requirements(self, dtype):
        if self.previous_controlnet is not None:
            return self.previous_controlnet.inference_memory_requirements(dtype)
        return 0

603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
    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

648
649
650
651
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
652
        self.control_model_wrapped = ModelPatcher(self.control_model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
653
        self.global_average_pooling = global_average_pooling
comfyanonymous's avatar
comfyanonymous committed
654

655
    def get_control(self, x_noisy, t, cond, batched_number):
comfyanonymous's avatar
comfyanonymous committed
656
657
        control_prev = None
        if self.previous_controlnet is not None:
658
            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
comfyanonymous's avatar
comfyanonymous committed
659

660
661
662
663
664
665
666
        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 {}

667
        output_dtype = x_noisy.dtype
comfyanonymous's avatar
comfyanonymous committed
668
669
670
671
        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
672
673
674
            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)
675
676


677
678
679
680
681
        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)
682
        return self.control_merge(None, control, control_prev, output_dtype)
comfyanonymous's avatar
comfyanonymous committed
683
684

    def copy(self):
685
        c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling)
686
        self.copy_to(c)
comfyanonymous's avatar
comfyanonymous committed
687
688
        return c

689
690
    def get_models(self):
        out = super().get_models()
comfyanonymous's avatar
comfyanonymous committed
691
        out.append(self.control_model_wrapped)
692
693
        return out

comfyanonymous's avatar
comfyanonymous committed
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
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
709
                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
710
            else:
comfyanonymous's avatar
comfyanonymous committed
711
                return torch.nn.functional.linear(input, self.weight.to(input.device), self.bias)
comfyanonymous's avatar
comfyanonymous committed
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747

    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
748
                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
749
            else:
comfyanonymous's avatar
comfyanonymous committed
750
                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
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771

    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)
772
773
        dtype = model.get_dtype()
        self.control_model.to(dtype)
comfyanonymous's avatar
comfyanonymous committed
774
775
776
777
778
779
        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
780
781
782
783
784
785
            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
786
            try:
comfyanonymous's avatar
comfyanonymous committed
787
                set_attr(self.control_model, k, weight)
comfyanonymous's avatar
comfyanonymous committed
788
789
790
791
792
            except:
                pass

        for k in self.control_weights:
            if k not in {"lora_controlnet"}:
793
                set_attr(self.control_model, k, self.control_weights[k].to(dtype).to(model_management.get_torch_device()))
comfyanonymous's avatar
comfyanonymous committed
794
795
796
797
798
799
800
801
802
803
804
805
806
807

    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
808

809
810
811
    def inference_memory_requirements(self, dtype):
        return utils.calculate_parameters(self.control_weights) * model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)

812
def load_controlnet(ckpt_path, model=None):
813
    controlnet_data = utils.load_torch_file(ckpt_path, safe_load=True)
comfyanonymous's avatar
comfyanonymous committed
814
815
    if "lora_controlnet" in controlnet_data:
        return ControlLora(controlnet_data)
816
817
818
819

    controlnet_config = None
    if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
        use_fp16 = model_management.should_use_fp16()
820
        controlnet_config = model_detection.unet_config_from_diffusers_unet(controlnet_data, use_fp16)
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
        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)

859
860
861
        leftover_keys = controlnet_data.keys()
        if len(leftover_keys) > 0:
            print("leftover keys:", leftover_keys)
862
863
        controlnet_data = new_sd

864
    pth_key = 'control_model.zero_convs.0.0.weight'
comfyanonymous's avatar
comfyanonymous committed
865
    pth = False
866
    key = 'zero_convs.0.0.weight'
comfyanonymous's avatar
comfyanonymous committed
867
868
869
    if pth_key in controlnet_data:
        pth = True
        key = pth_key
870
        prefix = "control_model."
comfyanonymous's avatar
comfyanonymous committed
871
    elif key in controlnet_data:
872
        prefix = ""
comfyanonymous's avatar
comfyanonymous committed
873
    else:
874
875
876
877
        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
878

879
880
881
    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
882
    controlnet_config.pop("out_channels")
883
    controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
884
885
    control_model = cldm.ControlNet(**controlnet_config)

comfyanonymous's avatar
comfyanonymous committed
886
    if pth:
887
888
        if 'difference' in controlnet_data:
            if model is not None:
889
890
                model_management.load_models_gpu([model])
                model_sd = model.model_state_dict()
891
892
893
                for x in controlnet_data:
                    c_m = "control_model."
                    if x.startswith(c_m):
comfyanonymous's avatar
comfyanonymous committed
894
                        sd_key = "diffusion_model.{}".format(x[len(c_m):])
895
896
897
898
899
900
                        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
901
902
903
904
        class WeightsLoader(torch.nn.Module):
            pass
        w = WeightsLoader()
        w.control_model = control_model
905
        missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
comfyanonymous's avatar
comfyanonymous committed
906
    else:
907
908
        missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
    print(missing, unexpected)
comfyanonymous's avatar
comfyanonymous committed
909

910
911
912
    if use_fp16:
        control_model = control_model.half()

913
914
915
916
917
    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
918
919
    return control

920
class T2IAdapter(ControlBase):
921
    def __init__(self, t2i_model, channels_in, device=None):
922
        super().__init__(device)
923
924
925
926
        self.t2i_model = t2i_model
        self.channels_in = channels_in
        self.control_input = None

927
928
929
930
931
932
    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

933
    def get_control(self, x_noisy, t, cond, batched_number):
934
935
        control_prev = None
        if self.previous_controlnet is not None:
936
            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
937

938
939
940
941
942
943
944
        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 {}

945
946
947
        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
948
            self.control_input = None
949
            self.cond_hint = None
950
951
            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)
952
953
            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
954
955
956
        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:
957
            self.t2i_model.to(x_noisy.dtype)
958
            self.t2i_model.to(self.device)
959
            self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
960
961
            self.t2i_model.cpu()

962
        control_input = list(map(lambda a: None if a is None else a.clone(), self.control_input))
comfyanonymous's avatar
comfyanonymous committed
963
964
965
966
967
        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)
968
969
970

    def copy(self):
        c = T2IAdapter(self.t2i_model, self.channels_in)
971
        self.copy_to(c)
972
973
        return c

974
def load_t2i_adapter(t2i_data):
975
    keys = t2i_data.keys()
976
977
978
    if 'adapter' in keys:
        t2i_data = t2i_data['adapter']
        keys = t2i_data.keys()
979
    if "body.0.in_conv.weight" in keys:
980
981
        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)
982
    elif 'conv_in.weight' in keys:
983
        cin = t2i_data['conv_in.weight'].shape[1]
984
985
986
987
988
989
        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
990
991
992
993
        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)
994
995
    else:
        return None
comfyanonymous's avatar
comfyanonymous committed
996
997
998
999
1000
1001
1002
1003
    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
1004

1005
1006
1007
1008
1009
1010
1011
1012
1013
1014

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):
1015
    model_data = utils.load_torch_file(ckpt_path, safe_load=True)
1016
1017
1018
1019
1020
1021
1022
1023
1024
    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)


1025
1026
1027
1028
1029
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
1030
1031
1032
    class EmptyClass:
        pass

1033
1034
1035
1036
    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
1037
1038
    clip_target = EmptyClass()
    clip_target.params = {}
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
    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
1049
    else:
1050
1051
        clip_target.clip = sdxl_clip.SDXLClipModel
        clip_target.tokenizer = sdxl_clip.SDXLTokenizer
comfyanonymous's avatar
comfyanonymous committed
1052
1053

    clip = CLIP(clip_target, embedding_directory=embedding_directory)
1054
1055
1056
1057
1058
1059
1060
    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)
1061
    return clip
comfyanonymous's avatar
comfyanonymous committed
1062

1063
def load_gligen(ckpt_path):
1064
    data = utils.load_torch_file(ckpt_path, safe_load=True)
1065
1066
1067
    model = gligen.load_gligen(data)
    if model_management.should_use_fp16():
        model = model.half()
comfyanonymous's avatar
comfyanonymous committed
1068
    return ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
1069

comfyanonymous's avatar
comfyanonymous committed
1070
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
1071
    #TODO: this function is a mess and should be removed eventually
comfyanonymous's avatar
comfyanonymous committed
1072
1073
1074
    if config is None:
        with open(config_path, 'r') as stream:
            config = yaml.safe_load(stream)
comfyanonymous's avatar
comfyanonymous committed
1075
1076
1077
1078
1079
    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']

1080
1081
1082
    fp16 = False
    if "unet_config" in model_config_params:
        if "params" in model_config_params["unet_config"]:
comfyanonymous's avatar
comfyanonymous committed
1083
1084
1085
1086
1087
1088
1089
1090
            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"]

1091
    model_type = model_base.ModelType.EPS
comfyanonymous's avatar
comfyanonymous committed
1092
1093
1094

    if "parameterization" in model_config_params:
        if model_config_params["parameterization"] == "v":
1095
            model_type = model_base.ModelType.V_PREDICTION
1096

comfyanonymous's avatar
comfyanonymous committed
1097
1098
1099
1100
1101
1102
    clip = None
    vae = None

    class WeightsLoader(torch.nn.Module):
        pass

1103
1104
    if state_dict is None:
        state_dict = utils.load_torch_file(ckpt_path)
comfyanonymous's avatar
comfyanonymous committed
1105

1106
1107
1108
1109
1110
1111
1112
1113
    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
1114
    if config['model']["target"].endswith("LatentInpaintDiffusion"):
1115
        model = model_base.SDInpaint(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1116
    elif config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
1117
        model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1118
    else:
1119
        model = model_base.BaseModel(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
1120

1121
1122
1123
    if fp16:
        model = model.half()

1124
1125
    offload_device = model_management.unet_offload_device()
    model = model.to(offload_device)
1126
1127
1128
1129
    model.load_model_weights(state_dict, "model.diffusion_model.")

    if output_vae:
        w = WeightsLoader()
1130
        vae = VAE(config=vae_config)
1131
1132
1133
1134
1135
1136
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, state_dict)

    if output_clip:
        w = WeightsLoader()
        clip_target = EmptyClass()
1137
        clip_target.params = clip_config.get("params", {})
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
        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)

1148
    return (ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
1149

1150
1151
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)
1152
1153
    sd_keys = sd.keys()
    clip = None
1154
    clipvision = None
1155
    vae = None
1156
1157
    model = None
    clip_target = None
1158

1159
    parameters = utils.calculate_parameters(sd, "model.diffusion_model.")
1160
    fp16 = model_management.should_use_fp16(model_params=parameters)
1161

1162
1163
1164
    class WeightsLoader(torch.nn.Module):
        pass

1165
1166
1167
    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))
1168

1169
    if model_config.clip_vision_prefix is not None:
1170
        if output_clipvision:
1171
            clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
1172

comfyanonymous's avatar
comfyanonymous committed
1173
1174
1175
1176
1177
    dtype = torch.float32
    if fp16:
        dtype = torch.float16

    inital_load_device = model_management.unet_inital_load_device(parameters, dtype)
1178
    offload_device = model_management.unet_offload_device()
comfyanonymous's avatar
comfyanonymous committed
1179
    model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
1180
    model.load_model_weights(sd, "model.diffusion_model.")
1181

1182
    if output_vae:
1183
        vae = VAE()
1184
1185
1186
        w = WeightsLoader()
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, sd)
1187

1188
1189
1190
1191
1192
1193
1194
    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
1195

1196
1197
1198
    left_over = sd.keys()
    if len(left_over) > 0:
        print("left over keys:", left_over)
1199

comfyanonymous's avatar
comfyanonymous committed
1200
1201
1202
1203
1204
1205
    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)
1206

1207
1208
1209

def load_unet(unet_path): #load unet in diffusers format
    sd = utils.load_torch_file(unet_path)
1210
    parameters = utils.calculate_parameters(sd)
1211
1212
    fp16 = model_management.should_use_fp16(model_params=parameters)

1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
    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)
1231

1232
def save_checkpoint(output_path, model, clip, vae, metadata=None):
1233
1234
1235
    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)