sd.py 30.3 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

12
13
import comfy.utils

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

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

24
import comfy.lora
25
import comfy.t2i_adapter.adapter
26

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
    sd = comfy.utils.transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.text_model.", 24)
54
    return load_model_weights(model, sd)
comfyanonymous's avatar
comfyanonymous committed
55

56
class ModelPatcher:
comfyanonymous's avatar
comfyanonymous committed
57
    def __init__(self, model, load_device, offload_device, size=0, current_device=None):
58
        self.size = size
59
        self.model = model
60
        self.patches = {}
61
        self.backup = {}
62
        self.model_options = {"transformer_options":{}}
63
        self.model_size()
64
65
        self.load_device = load_device
        self.offload_device = offload_device
comfyanonymous's avatar
comfyanonymous committed
66
67
68
69
        if current_device is None:
            self.current_device = self.offload_device
        else:
            self.current_device = current_device
70
71
72
73
74
75
76
77
78
79

    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
80
        self.model_keys = set(model_sd.keys())
81
        return size
82
83

    def clone(self):
comfyanonymous's avatar
comfyanonymous committed
84
        n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device)
85
86
87
88
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]

89
        n.model_options = copy.deepcopy(self.model_options)
90
        n.model_keys = self.model_keys
91
92
        return n

comfyanonymous's avatar
comfyanonymous committed
93
94
95
96
97
    def is_clone(self, other):
        if hasattr(other, 'model') and self.model is other.model:
            return True
        return False

98
    def set_model_sampler_cfg_function(self, sampler_cfg_function):
99
100
101
102
        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
103

104
105
106
    def set_model_unet_function_wrapper(self, unet_wrapper_function):
        self.model_options["model_function_wrapper"] = unet_wrapper_function

107
108
109
110
111
112
    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]

113
114
115
116
117
118
119
120
    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

121
122
123
124
125
126
    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")

127
128
129
130
131
132
133
134
135
    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")

136
137
138
    def set_model_attn2_output_patch(self, patch):
        self.set_model_patch(patch, "attn2_output_patch")

139
140
141
142
143
144
145
146
147
    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)
148
149
150
151
152
153
154
        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)
155

156
    def model_dtype(self):
157
158
        if hasattr(self.model, "get_dtype"):
            return self.model.get_dtype()
159

160
    def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
161
        p = set()
162
        for k in patches:
163
            if k in self.model_keys:
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
                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
183

184
    def model_state_dict(self, filter_prefix=None):
185
186
        sd = self.model.state_dict()
        keys = list(sd.keys())
187
188
189
190
        if filter_prefix is not None:
            for k in keys:
                if not k.startswith(filter_prefix):
                    sd.pop(k)
191
192
        return sd

193
    def patch_model(self, device_to=None):
194
        model_sd = self.model_state_dict()
195
196
197
198
        for key in self.patches:
            if key not in model_sd:
                print("could not patch. key doesn't exist in model:", k)
                continue
199

200
            weight = model_sd[key]
201

202
            if key not in self.backup:
203
                self.backup[key] = weight.to(self.offload_device)
204

205
206
207
208
            if device_to is not None:
                temp_weight = weight.float().to(device_to, copy=True)
            else:
                temp_weight = weight.to(torch.float32, copy=True)
209
            out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
210
            comfy.utils.set_attr(self.model, key, out_weight)
211
            del temp_weight
comfyanonymous's avatar
comfyanonymous committed
212
213
214
215
216

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

217
        return self.model
comfyanonymous's avatar
comfyanonymous committed
218

219
220
221
222
223
224
225
226
227
228
229
230
231
232
    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]
233
234
235
236
237
                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)
238
            elif len(v) == 4: #lora/locon
comfyanonymous's avatar
comfyanonymous committed
239
240
                mat1 = v[0].float().to(weight.device)
                mat2 = v[1].float().to(weight.device)
241
242
243
244
                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
245
246
247
                    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)
248
249
250
251
                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)
252
253
254
255
256
257
258
259
260
261
262
263
264
            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
265
266
                else:
                    w1 = w1.float().to(weight.device)
267
268
269
270

                if w2 is None:
                    dim = w2_b.shape[0]
                    if t2 is None:
comfyanonymous's avatar
comfyanonymous committed
271
                        w2 = torch.mm(w2_a.float().to(weight.device), w2_b.float().to(weight.device))
272
                    else:
comfyanonymous's avatar
comfyanonymous committed
273
274
275
                        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)
276
277
278
279
280
281

                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

282
283
284
285
                try:
                    weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)
286
287
288
289
290
291
292
293
294
295
            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
296
297
                    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))
298
                else:
comfyanonymous's avatar
comfyanonymous committed
299
300
                    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))
301

302
303
304
305
306
                try:
                    weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
                except Exception as e:
                    print("ERROR", key, e)

307
        return weight
308

comfyanonymous's avatar
comfyanonymous committed
309
    def unpatch_model(self, device_to=None):
310
        keys = list(self.backup.keys())
311

312
        for k in keys:
313
            comfy.utils.set_attr(self.model, k, self.backup[k])
314

315
316
        self.backup = {}

comfyanonymous's avatar
comfyanonymous committed
317
318
319
320
321
        if device_to is not None:
            self.model.to(device_to)
            self.current_device = device_to


322
def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
323
324
325
    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)
326
327
328
329
330
331
332
333
334
335
336
    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
337
338
339


class CLIP:
340
    def __init__(self, target=None, embedding_directory=None, no_init=False):
341
342
        if no_init:
            return
comfyanonymous's avatar
comfyanonymous committed
343
        params = target.params.copy()
344
345
        clip = target.clip
        tokenizer = target.tokenizer
346

347
348
        load_device = model_management.text_encoder_device()
        offload_device = model_management.text_encoder_offload_device()
comfyanonymous's avatar
comfyanonymous committed
349
        params['device'] = load_device
350
        if model_management.should_use_fp16(load_device, prioritize_performance=False):
351
352
353
354
355
            params['dtype'] = torch.float16
        else:
            params['dtype'] = torch.float32

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

357
        self.tokenizer = tokenizer(embedding_directory=embedding_directory)
358
        self.patcher = ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
359
        self.layer_idx = None
360
361
362
363
364
365

    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
366
        n.layer_idx = self.layer_idx
367
368
        return n

369
370
    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
371

372
    def clip_layer(self, layer_idx):
comfyanonymous's avatar
comfyanonymous committed
373
        self.layer_idx = layer_idx
374

375
376
    def tokenize(self, text, return_word_ids=False):
        return self.tokenizer.tokenize_with_weights(text, return_word_ids)
BlenderNeko's avatar
BlenderNeko committed
377

378
    def encode_from_tokens(self, tokens, return_pooled=False):
379
380
        if self.layer_idx is not None:
            self.cond_stage_model.clip_layer(self.layer_idx)
381
382
        else:
            self.cond_stage_model.reset_clip_layer()
383

384
        self.load_model()
385
        cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
386
        if return_pooled:
387
388
            return cond, pooled
        return cond
comfyanonymous's avatar
comfyanonymous committed
389

390
    def encode(self, text):
391
        tokens = self.tokenize(text)
392
393
        return self.encode_from_tokens(tokens)

394
395
    def load_sd(self, sd):
        return self.cond_stage_model.load_sd(sd)
396

397
398
399
    def get_sd(self):
        return self.cond_stage_model.state_dict()

400
401
402
    def load_model(self):
        model_management.load_model_gpu(self.patcher)
        return self.patcher
403

404
405
406
    def get_key_patches(self):
        return self.patcher.get_key_patches()

comfyanonymous's avatar
comfyanonymous committed
407
class VAE:
408
    def __init__(self, ckpt_path=None, device=None, config=None):
comfyanonymous's avatar
comfyanonymous committed
409
410
411
        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}
412
            self.first_stage_model = AutoencoderKL(ddconfig, {'target': 'torch.nn.Identity'}, 4, monitor="val/rec_loss")
comfyanonymous's avatar
comfyanonymous committed
413
        else:
414
            self.first_stage_model = AutoencoderKL(**(config['params']))
comfyanonymous's avatar
comfyanonymous committed
415
        self.first_stage_model = self.first_stage_model.eval()
416
        if ckpt_path is not None:
417
            sd = comfy.utils.load_torch_file(ckpt_path)
418
419
420
421
            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)

422
        if device is None:
423
            device = model_management.vae_device()
comfyanonymous's avatar
comfyanonymous committed
424
        self.device = device
425
        self.offload_device = model_management.vae_offload_device()
426
427
        self.vae_dtype = model_management.vae_dtype()
        self.first_stage_model.to(self.vae_dtype)
comfyanonymous's avatar
comfyanonymous committed
428

429
    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
430
431
432
433
        steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
        steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
        steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
        pbar = comfy.utils.ProgressBar(steps)
434

435
        decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
436
        output = torch.clamp((
437
438
439
            (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = 8, pbar = pbar) +
            comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = 8, pbar = pbar) +
             comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = 8, pbar = pbar))
440
441
442
            / 3.0) / 2.0, min=0.0, max=1.0)
        return output

443
    def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
444
445
446
447
        steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
        steps += pixel_samples.shape[0] * comfy.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] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)
        pbar = comfy.utils.ProgressBar(steps)
448

449
        encode_fn = lambda a: self.first_stage_model.encode(2. * a.to(self.vae_dtype).to(self.device) - 1.).sample().float()
450
451
452
        samples = comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
        samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
        samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/8), out_channels=4, pbar=pbar)
453
454
455
        samples /= 3.0
        return samples

456
    def decode(self, samples_in):
comfyanonymous's avatar
comfyanonymous committed
457
        self.first_stage_model = self.first_stage_model.to(self.device)
458
        try:
459
            memory_used = (2562 * samples_in.shape[2] * samples_in.shape[3] * 64) * 1.7
comfyanonymous's avatar
comfyanonymous committed
460
            model_management.free_memory(memory_used, self.device)
461
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
462
            batch_number = int(free_memory / memory_used)
463
464
465
466
            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):
467
468
                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()
469
470
471
472
        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)

473
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
474
475
476
        pixel_samples = pixel_samples.cpu().movedim(1,-1)
        return pixel_samples

477
    def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
478
        self.first_stage_model = self.first_stage_model.to(self.device)
479
        output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
480
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
481
482
        return output.movedim(1,-1)

comfyanonymous's avatar
comfyanonymous committed
483
484
    def encode(self, pixel_samples):
        self.first_stage_model = self.first_stage_model.to(self.device)
485
486
        pixel_samples = pixel_samples.movedim(-1,1)
        try:
487
            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
488
            model_management.free_memory(memory_used, self.device)
489
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
490
            batch_number = int(free_memory / memory_used)
491
            batch_number = max(1, batch_number)
492
493
            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):
494
495
                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()
496

497
498
499
500
        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)

501
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
502
503
        return samples

comfyanonymous's avatar
comfyanonymous committed
504
505
    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)
506
507
        pixel_samples = pixel_samples.movedim(-1,1)
        samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
508
        self.first_stage_model = self.first_stage_model.to(self.offload_device)
comfyanonymous's avatar
comfyanonymous committed
509
        return samples
510

511
512
513
    def get_sd(self):
        return self.first_stage_model.state_dict()

514
515
516
517
518
519
520
521
522
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):
523
    model_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
524
525
    keys = model_data.keys()
    if "style_embedding" in keys:
526
        model = comfy.t2i_adapter.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
527
528
529
530
531
532
    else:
        raise Exception("invalid style model {}".format(ckpt_path))
    model.load_state_dict(model_data)
    return StyleModel(model)


533
534
535
def load_clip(ckpt_paths, embedding_directory=None):
    clip_data = []
    for p in ckpt_paths:
536
        clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
537

comfyanonymous's avatar
comfyanonymous committed
538
539
540
    class EmptyClass:
        pass

541
542
    for i in range(len(clip_data)):
        if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
543
            clip_data[i] = comfy.utils.transformers_convert(clip_data[i], "", "text_model.", 32)
544

comfyanonymous's avatar
comfyanonymous committed
545
546
    clip_target = EmptyClass()
    clip_target.params = {}
547
548
549
550
551
552
553
554
555
556
    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
557
    else:
558
559
        clip_target.clip = sdxl_clip.SDXLClipModel
        clip_target.tokenizer = sdxl_clip.SDXLTokenizer
comfyanonymous's avatar
comfyanonymous committed
560
561

    clip = CLIP(clip_target, embedding_directory=embedding_directory)
562
563
564
565
566
567
568
    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)
569
    return clip
comfyanonymous's avatar
comfyanonymous committed
570

571
def load_gligen(ckpt_path):
572
    data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
573
574
575
    model = gligen.load_gligen(data)
    if model_management.should_use_fp16():
        model = model.half()
comfyanonymous's avatar
comfyanonymous committed
576
    return ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
577

comfyanonymous's avatar
comfyanonymous committed
578
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
579
    #TODO: this function is a mess and should be removed eventually
comfyanonymous's avatar
comfyanonymous committed
580
581
582
    if config is None:
        with open(config_path, 'r') as stream:
            config = yaml.safe_load(stream)
comfyanonymous's avatar
comfyanonymous committed
583
584
585
586
587
    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']

588
589
590
    fp16 = False
    if "unet_config" in model_config_params:
        if "params" in model_config_params["unet_config"]:
comfyanonymous's avatar
comfyanonymous committed
591
592
593
594
595
596
597
598
            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"]

599
    model_type = model_base.ModelType.EPS
comfyanonymous's avatar
comfyanonymous committed
600
601
602

    if "parameterization" in model_config_params:
        if model_config_params["parameterization"] == "v":
603
            model_type = model_base.ModelType.V_PREDICTION
604

comfyanonymous's avatar
comfyanonymous committed
605
606
607
608
609
610
    clip = None
    vae = None

    class WeightsLoader(torch.nn.Module):
        pass

611
    if state_dict is None:
612
        state_dict = comfy.utils.load_torch_file(ckpt_path)
comfyanonymous's avatar
comfyanonymous committed
613

614
615
616
617
618
619
620
621
    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
622
    if config['model']["target"].endswith("LatentInpaintDiffusion"):
623
        model = model_base.SDInpaint(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
624
    elif config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
625
        model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
626
    else:
627
        model = model_base.BaseModel(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
628

629
630
631
    if fp16:
        model = model.half()

632
633
    offload_device = model_management.unet_offload_device()
    model = model.to(offload_device)
634
635
636
637
    model.load_model_weights(state_dict, "model.diffusion_model.")

    if output_vae:
        w = WeightsLoader()
638
        vae = VAE(config=vae_config)
639
640
641
642
643
644
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, state_dict)

    if output_clip:
        w = WeightsLoader()
        clip_target = EmptyClass()
645
        clip_target.params = clip_config.get("params", {})
646
647
648
649
650
651
652
653
654
655
        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)

656
    return (ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
657

658
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None):
659
    sd = comfy.utils.load_torch_file(ckpt_path)
660
661
    sd_keys = sd.keys()
    clip = None
662
    clipvision = None
663
    vae = None
664
665
    model = None
    clip_target = None
666

667
    parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.")
668
    fp16 = model_management.should_use_fp16(model_params=parameters)
669

670
671
672
    class WeightsLoader(torch.nn.Module):
        pass

673
674
675
    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))
676

677
    if model_config.clip_vision_prefix is not None:
678
        if output_clipvision:
679
            clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
680

comfyanonymous's avatar
comfyanonymous committed
681
682
683
684
685
    dtype = torch.float32
    if fp16:
        dtype = torch.float16

    inital_load_device = model_management.unet_inital_load_device(parameters, dtype)
686
    offload_device = model_management.unet_offload_device()
comfyanonymous's avatar
comfyanonymous committed
687
    model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
688
    model.load_model_weights(sd, "model.diffusion_model.")
689

690
    if output_vae:
691
        vae = VAE()
692
693
694
        w = WeightsLoader()
        w.first_stage_model = vae.first_stage_model
        load_model_weights(w, sd)
695

696
697
698
699
700
701
702
    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
703

704
705
706
    left_over = sd.keys()
    if len(left_over) > 0:
        print("left over keys:", left_over)
707

comfyanonymous's avatar
comfyanonymous committed
708
709
710
711
712
713
    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)
714

715
716

def load_unet(unet_path): #load unet in diffusers format
717
718
    sd = comfy.utils.load_torch_file(unet_path)
    parameters = comfy.utils.calculate_parameters(sd)
719
720
    fp16 = model_management.should_use_fp16(model_params=parameters)

721
722
723
724
725
    model_config = model_detection.model_config_from_diffusers_unet(sd, fp16)
    if model_config is None:
        print("ERROR UNSUPPORTED UNET", unet_path)
        return None

726
    diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)
727
728
729
730
731
732
733
734
735
736
737
738

    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)
739

740
def save_checkpoint(output_path, model, clip, vae, metadata=None):
741
742
    model_management.load_models_gpu([model, clip.load_model()])
    sd = model.model.state_dict_for_saving(clip.get_sd(), vae.get_sd())
743
    comfy.utils.save_torch_file(sd, output_path, metadata=metadata)