sd.py 27.1 KB
Newer Older
comfyanonymous's avatar
comfyanonymous committed
1
import torch
2
from enum import Enum
comfyanonymous's avatar
comfyanonymous committed
3

4
from comfy import model_management
comfyanonymous's avatar
comfyanonymous committed
5
from .ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
comfyanonymous's avatar
comfyanonymous committed
6
from .ldm.cascade.stage_a import StageA
7
from .ldm.cascade.stage_c_coder import StageC_coder
comfyanonymous's avatar
comfyanonymous committed
8

9
import yaml
comfyanonymous's avatar
comfyanonymous committed
10

11
12
import comfy.utils

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

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

23
import comfy.model_patcher
24
import comfy.lora
25
import comfy.t2i_adapter.adapter
26
import comfy.supported_models_base
27
import comfy.taesd.taesd
28

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

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

55
    sd = comfy.utils.clip_text_transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.")
56
    return load_model_weights(model, sd)
comfyanonymous's avatar
comfyanonymous committed
57

comfyanonymous's avatar
comfyanonymous committed
58

59
def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
60
61
62
63
64
65
    key_map = {}
    if model is not None:
        key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
    if clip is not None:
        key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)

66
    loaded = comfy.lora.load_lora(lora, key_map)
67
68
69
70
71
72
73
74
75
76
77
78
79
    if model is not None:
        new_modelpatcher = model.clone()
        k = new_modelpatcher.add_patches(loaded, strength_model)
    else:
        k = ()
        new_modelpatcher = None

    if clip is not None:
        new_clip = clip.clone()
        k1 = new_clip.add_patches(loaded, strength_clip)
    else:
        k1 = ()
        new_clip = None
80
81
82
83
84
85
86
    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
87
88
89


class CLIP:
90
    def __init__(self, target=None, embedding_directory=None, no_init=False):
91
92
        if no_init:
            return
comfyanonymous's avatar
comfyanonymous committed
93
        params = target.params.copy()
94
95
        clip = target.clip
        tokenizer = target.tokenizer
96

97
98
        load_device = model_management.text_encoder_device()
        offload_device = model_management.text_encoder_offload_device()
99
        params['device'] = offload_device
100
        params['dtype'] = model_management.text_encoder_dtype(load_device)
101
102

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

104
        self.tokenizer = tokenizer(embedding_directory=embedding_directory)
105
        self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
106
        self.layer_idx = None
107
108
109
110
111
112

    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
113
        n.layer_idx = self.layer_idx
114
115
        return n

116
117
    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
118

119
    def clip_layer(self, layer_idx):
comfyanonymous's avatar
comfyanonymous committed
120
        self.layer_idx = layer_idx
121

122
123
    def tokenize(self, text, return_word_ids=False):
        return self.tokenizer.tokenize_with_weights(text, return_word_ids)
BlenderNeko's avatar
BlenderNeko committed
124

125
    def encode_from_tokens(self, tokens, return_pooled=False):
126
127
        self.cond_stage_model.reset_clip_options()

128
        if self.layer_idx is not None:
129
130
131
132
            self.cond_stage_model.set_clip_options({"layer": self.layer_idx})

        if return_pooled == "unprojected":
            self.cond_stage_model.set_clip_options({"projected_pooled": False})
133

134
        self.load_model()
135
        cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
136
        if return_pooled:
137
138
            return cond, pooled
        return cond
comfyanonymous's avatar
comfyanonymous committed
139

140
    def encode(self, text):
141
        tokens = self.tokenize(text)
142
143
        return self.encode_from_tokens(tokens)

144
145
146
147
148
    def load_sd(self, sd, full_model=False):
        if full_model:
            return self.cond_stage_model.load_state_dict(sd, strict=False)
        else:
            return self.cond_stage_model.load_sd(sd)
149

150
151
152
    def get_sd(self):
        return self.cond_stage_model.state_dict()

153
154
155
    def load_model(self):
        model_management.load_model_gpu(self.patcher)
        return self.patcher
156

157
158
159
    def get_key_patches(self):
        return self.patcher.get_key_patches()

comfyanonymous's avatar
comfyanonymous committed
160
class VAE:
161
    def __init__(self, sd=None, device=None, config=None, dtype=None):
comfyanonymous's avatar
comfyanonymous committed
162
163
164
        if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
            sd = diffusers_convert.convert_vae_state_dict(sd)

165
166
        self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower)
        self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype)
comfyanonymous's avatar
comfyanonymous committed
167
        self.downscale_ratio = 8
168
        self.upscale_ratio = 8
comfyanonymous's avatar
comfyanonymous committed
169
        self.latent_channels = 4
comfyanonymous's avatar
comfyanonymous committed
170
171
        self.process_input = lambda image: image * 2.0 - 1.0
        self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
172

comfyanonymous's avatar
comfyanonymous committed
173
        if config is None:
comfyanonymous's avatar
comfyanonymous committed
174
175
176
177
178
179
180
181
182
            if "decoder.mid.block_1.mix_factor" in sd:
                encoder_config = {'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}
                decoder_config = encoder_config.copy()
                decoder_config["video_kernel_size"] = [3, 1, 1]
                decoder_config["alpha"] = 0.0
                self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
                                                            encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config},
                                                            decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config})
            elif "taesd_decoder.1.weight" in sd:
183
                self.first_stage_model = comfy.taesd.taesd.TAESD()
comfyanonymous's avatar
comfyanonymous committed
184
185
186
            elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade
                self.first_stage_model = StageA()
                self.downscale_ratio = 4
187
                self.upscale_ratio = 4
comfyanonymous's avatar
comfyanonymous committed
188
189
190
191
192
                #TODO
                #self.memory_used_encode
                #self.memory_used_decode
                self.process_input = lambda image: image
                self.process_output = lambda image: image
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
            elif "backbone.1.0.block.0.1.num_batches_tracked" in sd: #effnet: encoder for stage c latent of stable cascade
                self.first_stage_model = StageC_coder()
                self.downscale_ratio = 32
                self.latent_channels = 16
                new_sd = {}
                for k in sd:
                    new_sd["encoder.{}".format(k)] = sd[k]
                sd = new_sd
            elif "blocks.11.num_batches_tracked" in sd: #previewer: decoder for stage c latent of stable cascade
                self.first_stage_model = StageC_coder()
                self.latent_channels = 16
                new_sd = {}
                for k in sd:
                    new_sd["previewer.{}".format(k)] = sd[k]
                sd = new_sd
            elif "encoder.backbone.1.0.block.0.1.num_batches_tracked" in sd: #combined effnet and previewer for stable cascade
                self.first_stage_model = StageC_coder()
                self.downscale_ratio = 32
                self.latent_channels = 16
212
213
214
            else:
                #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}
215
216
217
218

                if 'encoder.down.2.downsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
                    ddconfig['ch_mult'] = [1, 2, 4]
                    self.downscale_ratio = 4
219
                    self.upscale_ratio = 4
220

221
                self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=4)
comfyanonymous's avatar
comfyanonymous committed
222
        else:
223
            self.first_stage_model = AutoencoderKL(**(config['params']))
comfyanonymous's avatar
comfyanonymous committed
224
        self.first_stage_model = self.first_stage_model.eval()
comfyanonymous's avatar
comfyanonymous committed
225
226
227
228
229
230
231

        m, u = self.first_stage_model.load_state_dict(sd, strict=False)
        if len(m) > 0:
            print("Missing VAE keys", m)

        if len(u) > 0:
            print("Leftover VAE keys", u)
232

233
        if device is None:
234
            device = model_management.vae_device()
comfyanonymous's avatar
comfyanonymous committed
235
        self.device = device
236
        offload_device = model_management.vae_offload_device()
237
238
239
        if dtype is None:
            dtype = model_management.vae_dtype()
        self.vae_dtype = dtype
240
        self.first_stage_model.to(self.vae_dtype)
241
        self.output_device = model_management.intermediate_device()
comfyanonymous's avatar
comfyanonymous committed
242

243
244
        self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)

245
246
247
248
249
250
251
252
253
    def vae_encode_crop_pixels(self, pixels):
        x = (pixels.shape[1] // self.downscale_ratio) * self.downscale_ratio
        y = (pixels.shape[2] // self.downscale_ratio) * self.downscale_ratio
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % self.downscale_ratio) // 2
            y_offset = (pixels.shape[2] % self.downscale_ratio) // 2
            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
        return pixels

254
    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
255
256
257
258
        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)
259

comfyanonymous's avatar
comfyanonymous committed
260
261
        decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
        output = self.process_output(
262
263
264
            (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
            comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
             comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar))
comfyanonymous's avatar
comfyanonymous committed
265
            / 3.0)
266
267
        return output

268
    def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
269
270
271
272
        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)
273

comfyanonymous's avatar
comfyanonymous committed
274
        encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()
comfyanonymous's avatar
comfyanonymous committed
275
276
277
        samples = comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
        samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
        samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
278
279
280
        samples /= 3.0
        return samples

281
282
    def decode(self, samples_in):
        try:
283
            memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
284
            model_management.load_models_gpu([self.patcher], memory_required=memory_used)
285
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
286
            batch_number = int(free_memory / memory_used)
287
288
            batch_number = max(1, batch_number)

289
            pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.upscale_ratio), round(samples_in.shape[3] * self.upscale_ratio)), device=self.output_device)
290
            for x in range(0, samples_in.shape[0], batch_number):
291
                samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
comfyanonymous's avatar
comfyanonymous committed
292
                pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())
293
294
295
296
        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)

297
        pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
comfyanonymous's avatar
comfyanonymous committed
298
299
        return pixel_samples

300
    def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
301
        model_management.load_model_gpu(self.patcher)
302
        output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
303
304
        return output.movedim(1,-1)

comfyanonymous's avatar
comfyanonymous committed
305
    def encode(self, pixel_samples):
306
        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
307
308
        pixel_samples = pixel_samples.movedim(-1,1)
        try:
309
            memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
310
            model_management.load_models_gpu([self.patcher], memory_required=memory_used)
311
            free_memory = model_management.get_free_memory(self.device)
comfyanonymous's avatar
comfyanonymous committed
312
            batch_number = int(free_memory / memory_used)
313
            batch_number = max(1, batch_number)
comfyanonymous's avatar
comfyanonymous committed
314
            samples = torch.empty((pixel_samples.shape[0], self.latent_channels, round(pixel_samples.shape[2] // self.downscale_ratio), round(pixel_samples.shape[3] // self.downscale_ratio)), device=self.output_device)
315
            for x in range(0, pixel_samples.shape[0], batch_number):
comfyanonymous's avatar
comfyanonymous committed
316
                pixels_in = self.process_input(pixel_samples[x:x+batch_number]).to(self.vae_dtype).to(self.device)
317
                samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).to(self.output_device).float()
318

319
320
321
322
        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)

comfyanonymous's avatar
comfyanonymous committed
323
324
        return samples

comfyanonymous's avatar
comfyanonymous committed
325
    def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
326
        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
327
        model_management.load_model_gpu(self.patcher)
328
329
        pixel_samples = pixel_samples.movedim(-1,1)
        samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
comfyanonymous's avatar
comfyanonymous committed
330
        return samples
331

332
333
334
    def get_sd(self):
        return self.first_stage_model.state_dict()

335
336
337
338
339
340
341
342
343
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):
344
    model_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
345
346
    keys = model_data.keys()
    if "style_embedding" in keys:
347
        model = comfy.t2i_adapter.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
348
349
350
351
352
    else:
        raise Exception("invalid style model {}".format(ckpt_path))
    model.load_state_dict(model_data)
    return StyleModel(model)

353
354
355
class CLIPType(Enum):
    STABLE_DIFFUSION = 1
    STABLE_CASCADE = 2
356

357
def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION):
358
359
    clip_data = []
    for p in ckpt_paths:
360
        clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
361

comfyanonymous's avatar
comfyanonymous committed
362
363
364
    class EmptyClass:
        pass

365
366
    for i in range(len(clip_data)):
        if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
367
            clip_data[i] = comfy.utils.clip_text_transformers_convert(clip_data[i], "", "")
368

comfyanonymous's avatar
comfyanonymous committed
369
370
    clip_target = EmptyClass()
    clip_target.params = {}
371
372
    if len(clip_data) == 1:
        if "text_model.encoder.layers.30.mlp.fc1.weight" in clip_data[0]:
373
374
375
376
377
378
            if clip_type == CLIPType.STABLE_CASCADE:
                clip_target.clip = sdxl_clip.StableCascadeClipModel
                clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer
            else:
                clip_target.clip = sdxl_clip.SDXLRefinerClipModel
                clip_target.tokenizer = sdxl_clip.SDXLTokenizer
379
380
381
382
383
384
        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
385
    else:
386
387
        clip_target.clip = sdxl_clip.SDXLClipModel
        clip_target.tokenizer = sdxl_clip.SDXLTokenizer
comfyanonymous's avatar
comfyanonymous committed
388
389

    clip = CLIP(clip_target, embedding_directory=embedding_directory)
390
391
392
393
394
395
396
    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)
397
    return clip
comfyanonymous's avatar
comfyanonymous committed
398

399
def load_gligen(ckpt_path):
400
    data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
401
402
403
    model = gligen.load_gligen(data)
    if model_management.should_use_fp16():
        model = model.half()
404
    return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
405

comfyanonymous's avatar
comfyanonymous committed
406
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
407
    #TODO: this function is a mess and should be removed eventually
comfyanonymous's avatar
comfyanonymous committed
408
409
410
    if config is None:
        with open(config_path, 'r') as stream:
            config = yaml.safe_load(stream)
comfyanonymous's avatar
comfyanonymous committed
411
412
413
414
415
    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']

416
417
418
    fp16 = False
    if "unet_config" in model_config_params:
        if "params" in model_config_params["unet_config"]:
comfyanonymous's avatar
comfyanonymous committed
419
420
            unet_config = model_config_params["unet_config"]["params"]
            if "use_fp16" in unet_config:
421
422
423
                fp16 = unet_config.pop("use_fp16")
                if fp16:
                    unet_config["dtype"] = torch.float16
comfyanonymous's avatar
comfyanonymous committed
424
425
426
427
428

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

429
    model_type = model_base.ModelType.EPS
comfyanonymous's avatar
comfyanonymous committed
430
431
432

    if "parameterization" in model_config_params:
        if model_config_params["parameterization"] == "v":
433
            model_type = model_base.ModelType.V_PREDICTION
434

comfyanonymous's avatar
comfyanonymous committed
435
436
437
438
439
440
    clip = None
    vae = None

    class WeightsLoader(torch.nn.Module):
        pass

441
    if state_dict is None:
442
        state_dict = comfy.utils.load_torch_file(ckpt_path)
comfyanonymous's avatar
comfyanonymous committed
443

444
445
446
    class EmptyClass:
        pass

447
448
    model_config = comfy.supported_models_base.BASE({})

449
450
    from . import latent_formats
    model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)
451
    model_config.unet_config = model_detection.convert_config(unet_config)
452

comfyanonymous's avatar
comfyanonymous committed
453
    if config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
454
        model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
455
    else:
456
        model = model_base.BaseModel(model_config, model_type=model_type)
comfyanonymous's avatar
comfyanonymous committed
457

comfyanonymous's avatar
comfyanonymous committed
458
459
460
    if config['model']["target"].endswith("LatentInpaintDiffusion"):
        model.set_inpaint()

461
462
463
    if fp16:
        model = model.half()

464
465
    offload_device = model_management.unet_offload_device()
    model = model.to(offload_device)
466
467
468
    model.load_model_weights(state_dict, "model.diffusion_model.")

    if output_vae:
comfyanonymous's avatar
comfyanonymous committed
469
470
        vae_sd = comfy.utils.state_dict_prefix_replace(state_dict, {"first_stage_model.": ""}, filter_keys=True)
        vae = VAE(sd=vae_sd, config=vae_config)
471
472
473
474

    if output_clip:
        w = WeightsLoader()
        clip_target = EmptyClass()
475
        clip_target.params = clip_config.get("params", {})
476
477
478
        if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
            clip_target.clip = sd2_clip.SD2ClipModel
            clip_target.tokenizer = sd2_clip.SD2Tokenizer
479
480
            clip = CLIP(clip_target, embedding_directory=embedding_directory)
            w.cond_stage_model = clip.cond_stage_model.clip_h
481
482
483
        elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
            clip_target.clip = sd1_clip.SD1ClipModel
            clip_target.tokenizer = sd1_clip.SD1Tokenizer
484
485
            clip = CLIP(clip_target, embedding_directory=embedding_directory)
            w.cond_stage_model = clip.cond_stage_model.clip_l
486
487
        load_clip_weights(w, state_dict)

488
    return (comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
489

490
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):
491
    sd = comfy.utils.load_torch_file(ckpt_path)
492
493
    sd_keys = sd.keys()
    clip = None
494
    clipvision = None
495
    vae = None
496
    model = None
497
    model_patcher = None
498
    clip_target = None
499

500
    parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.")
501
    load_device = model_management.get_torch_device()
502

comfyanonymous's avatar
comfyanonymous committed
503
504
505
506
    model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.")
    unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
507

508
509
    if model_config is None:
        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
510

511
    if model_config.clip_vision_prefix is not None:
512
        if output_clipvision:
513
            clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
514

515
    if output_model:
516
        inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)
517
518
519
        offload_device = model_management.unet_offload_device()
        model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
        model.load_model_weights(sd, "model.diffusion_model.")
520

521
    if output_vae:
522
        vae_sd = comfy.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)
523
        vae_sd = model_config.process_vae_state_dict(vae_sd)
comfyanonymous's avatar
comfyanonymous committed
524
        vae = VAE(sd=vae_sd)
525

526
527
    if output_clip:
        clip_target = model_config.clip_target()
comfyanonymous's avatar
comfyanonymous committed
528
        if clip_target is not None:
529
530
            clip_sd = model_config.process_clip_state_dict(sd)
            if len(clip_sd) > 0:
531
                clip = CLIP(clip_target, embedding_directory=embedding_directory)
532
533
534
535
536
537
                m, u = clip.load_sd(clip_sd, full_model=True)
                if len(m) > 0:
                    print("clip missing:", m)

                if len(u) > 0:
                    print("clip unexpected:", u)
538
539
            else:
                print("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
comfyanonymous's avatar
comfyanonymous committed
540

541
542
543
    left_over = sd.keys()
    if len(left_over) > 0:
        print("left over keys:", left_over)
544

545
    if output_model:
546
        model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
547
548
549
        if inital_load_device != torch.device("cpu"):
            print("loaded straight to GPU")
            model_management.load_model_gpu(model_patcher)
comfyanonymous's avatar
comfyanonymous committed
550
551

    return (model_patcher, clip, vae, clipvision)
552

553

554
def load_unet_state_dict(sd): #load unet in diffusers format
555
    parameters = comfy.utils.calculate_parameters(sd)
556
    unet_dtype = model_management.unet_dtype(model_params=parameters)
557
558
    load_device = model_management.get_torch_device()

comfyanonymous's avatar
comfyanonymous committed
559
560
    if "input_blocks.0.0.weight" in sd or 'clf.1.weight' in sd: #ldm or stable cascade
        model_config = model_detection.model_config_from_unet(sd, "")
561
        if model_config is None:
562
            return None
563
564
565
        new_sd = sd

    else: #diffusers
comfyanonymous's avatar
comfyanonymous committed
566
        model_config = model_detection.model_config_from_diffusers_unet(sd)
567
568
569
570
571
572
573
574
575
576
577
        if model_config is None:
            return None

        diffusers_keys = comfy.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)
comfyanonymous's avatar
comfyanonymous committed
578

579
    offload_device = model_management.unet_offload_device()
comfyanonymous's avatar
comfyanonymous committed
580
581
582
    unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
583
584
585
    model = model_config.get_model(new_sd, "")
    model = model.to(offload_device)
    model.load_model_weights(new_sd, "")
586
587
588
    left_over = sd.keys()
    if len(left_over) > 0:
        print("left over keys in unet:", left_over)
589
    return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
590

591
592
593
594
595
596
597
598
def load_unet(unet_path):
    sd = comfy.utils.load_torch_file(unet_path)
    model = load_unet_state_dict(sd)
    if model is None:
        print("ERROR UNSUPPORTED UNET", unet_path)
        raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
    return model

599
600
601
602
603
604
605
606
607
608
def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None):
    clip_sd = None
    load_models = [model]
    if clip is not None:
        load_models.append(clip.load_model())
        clip_sd = clip.get_sd()

    model_management.load_models_gpu(load_models)
    clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None
    sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), clip_vision_sd)
609
    comfy.utils.save_torch_file(sd, output_path, metadata=metadata)