nodes.py 57.4 KB
Newer Older
comfyanonymous's avatar
comfyanonymous committed
1
2
3
4
5
import torch

import os
import sys
import json
6
import hashlib
7
import traceback
8
import math
9
import time
comfyanonymous's avatar
comfyanonymous committed
10

11
from PIL import Image, ImageOps
comfyanonymous's avatar
comfyanonymous committed
12
13
from PIL.PngImagePlugin import PngInfo
import numpy as np
14
import safetensors.torch
comfyanonymous's avatar
comfyanonymous committed
15

comfyanonymous's avatar
comfyanonymous committed
16
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
comfyanonymous's avatar
comfyanonymous committed
17
18


19
import comfy.diffusers_load
comfyanonymous's avatar
comfyanonymous committed
20
import comfy.samplers
21
import comfy.sample
comfyanonymous's avatar
comfyanonymous committed
22
import comfy.sd
comfyanonymous's avatar
comfyanonymous committed
23
24
import comfy.utils

25
import comfy.clip_vision
26

27
import comfy.model_management
28
import importlib
comfyanonymous's avatar
comfyanonymous committed
29

30
import folder_paths
31
import latent_preview
space-nuko's avatar
space-nuko committed
32

33
def before_node_execution():
34
    comfy.model_management.throw_exception_if_processing_interrupted()
35

36
def interrupt_processing(value=True):
37
    comfy.model_management.interrupt_current_processing(value)
38

39
40
MAX_RESOLUTION=8192

comfyanonymous's avatar
comfyanonymous committed
41
42
43
class CLIPTextEncode:
    @classmethod
    def INPUT_TYPES(s):
44
        return {"required": {"text": ("STRING", {"multiline": True}), "clip": ("CLIP", )}}
comfyanonymous's avatar
comfyanonymous committed
45
46
47
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"

48
49
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
50
    def encode(self, clip, text):
51
52
53
        tokens = clip.tokenize(text)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        return ([[cond, {"pooled_output": pooled}]], )
comfyanonymous's avatar
comfyanonymous committed
54
55
56
57
58
59
60
61

class ConditioningCombine:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning_1": ("CONDITIONING", ), "conditioning_2": ("CONDITIONING", )}}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "combine"

62
63
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
64
65
66
    def combine(self, conditioning_1, conditioning_2):
        return (conditioning_1 + conditioning_2, )

FizzleDorf's avatar
FizzleDorf committed
67
68
69
class ConditioningAverage :
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
70
71
        return {"required": {"conditioning_to": ("CONDITIONING", ), "conditioning_from": ("CONDITIONING", ),
                              "conditioning_to_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
FizzleDorf's avatar
FizzleDorf committed
72
73
74
75
76
77
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "addWeighted"

    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
78
    def addWeighted(self, conditioning_to, conditioning_from, conditioning_to_strength):
FizzleDorf's avatar
FizzleDorf committed
79
        out = []
comfyanonymous's avatar
comfyanonymous committed
80
81
82
83
84

        if len(conditioning_from) > 1:
            print("Warning: ConditioningAverage conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")

        cond_from = conditioning_from[0][0]
85
        pooled_output_from = conditioning_from[0][1].get("pooled_output", None)
comfyanonymous's avatar
comfyanonymous committed
86
87
88

        for i in range(len(conditioning_to)):
            t1 = conditioning_to[i][0]
89
            pooled_output_to = conditioning_to[i][1].get("pooled_output", pooled_output_from)
comfyanonymous's avatar
comfyanonymous committed
90
91
92
93
94
            t0 = cond_from[:,:t1.shape[1]]
            if t0.shape[1] < t1.shape[1]:
                t0 = torch.cat([t0] + [torch.zeros((1, (t1.shape[1] - t0.shape[1]), t1.shape[2]))], dim=1)

            tw = torch.mul(t1, conditioning_to_strength) + torch.mul(t0, (1.0 - conditioning_to_strength))
95
96
97
98
99
100
101
            t_to = conditioning_to[i][1].copy()
            if pooled_output_from is not None and pooled_output_to is not None:
                t_to["pooled_output"] = torch.mul(pooled_output_to, conditioning_to_strength) + torch.mul(pooled_output_from, (1.0 - conditioning_to_strength))
            elif pooled_output_from is not None:
                t_to["pooled_output"] = pooled_output_from

            n = [tw, t_to]
FizzleDorf's avatar
FizzleDorf committed
102
103
104
            out.append(n)
        return (out, )

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
class ConditioningConcat:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "conditioning_to": ("CONDITIONING",),
            "conditioning_from": ("CONDITIONING",),
            }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "concat"

    CATEGORY = "advanced/conditioning"

    def concat(self, conditioning_to, conditioning_from):
        out = []

        if len(conditioning_from) > 1:
            print("Warning: ConditioningConcat conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")

        cond_from = conditioning_from[0][0]

        for i in range(len(conditioning_to)):
            t1 = conditioning_to[i][0]
            tw = torch.cat((t1, cond_from),1)
            n = [tw, conditioning_to[i][1].copy()]
            out.append(n)

        return (out, )

comfyanonymous's avatar
comfyanonymous committed
133
134
135
136
class ConditioningSetArea:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
137
138
139
140
                              "width": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
comfyanonymous's avatar
comfyanonymous committed
141
142
143
144
145
                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "append"

146
147
    CATEGORY = "conditioning"

148
    def append(self, conditioning, width, height, x, y, strength):
comfyanonymous's avatar
comfyanonymous committed
149
150
151
152
153
        c = []
        for t in conditioning:
            n = [t[0], t[1].copy()]
            n[1]['area'] = (height // 8, width // 8, y // 8, x // 8)
            n[1]['strength'] = strength
154
            n[1]['set_area_to_bounds'] = False
comfyanonymous's avatar
comfyanonymous committed
155
            c.append(n)
comfyanonymous's avatar
comfyanonymous committed
156
        return (c, )
comfyanonymous's avatar
comfyanonymous committed
157

Jacob Segal's avatar
Jacob Segal committed
158
159
160
161
162
163
class ConditioningSetMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
                              "mask": ("MASK", ),
                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
164
                              "set_cond_area": (["default", "mask bounds"],),
Jacob Segal's avatar
Jacob Segal committed
165
166
167
168
169
170
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "append"

    CATEGORY = "conditioning"

171
    def append(self, conditioning, mask, set_cond_area, strength):
Jacob Segal's avatar
Jacob Segal committed
172
        c = []
173
174
175
        set_area_to_bounds = False
        if set_cond_area != "default":
            set_area_to_bounds = True
Jacob Segal's avatar
Jacob Segal committed
176
177
178
179
180
181
        if len(mask.shape) < 3:
            mask = mask.unsqueeze(0)
        for t in conditioning:
            n = [t[0], t[1].copy()]
            _, h, w = mask.shape
            n[1]['mask'] = mask
Jacob Segal's avatar
Jacob Segal committed
182
            n[1]['set_area_to_bounds'] = set_area_to_bounds
183
            n[1]['mask_strength'] = strength
Jacob Segal's avatar
Jacob Segal committed
184
185
186
            c.append(n)
        return (c, )

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
class ConditioningZeroOut:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", )}}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "zero_out"

    CATEGORY = "advanced/conditioning"

    def zero_out(self, conditioning):
        c = []
        for t in conditioning:
            d = t[1].copy()
            if "pooled_output" in d:
                d["pooled_output"] = torch.zeros_like(d["pooled_output"])
            n = [torch.zeros_like(t[0]), d]
            c.append(n)
        return (c, )

comfyanonymous's avatar
comfyanonymous committed
206
207
208
209
210
211
212
class VAEDecode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"

213
214
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
215
    def decode(self, vae, samples):
216
        return (vae.decode(samples["samples"]), )
comfyanonymous's avatar
comfyanonymous committed
217

218
219
220
221
222
223
224
225
226
227
228
229
class VAEDecodeTiled:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"

    CATEGORY = "_for_testing"

    def decode(self, vae, samples):
        return (vae.decode_tiled(samples["samples"]), )

comfyanonymous's avatar
comfyanonymous committed
230
231
232
233
234
235
236
class VAEEncode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

237
238
    CATEGORY = "latent"

239
240
241
242
    @staticmethod
    def vae_encode_crop_pixels(pixels):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
243
        if pixels.shape[1] != x or pixels.shape[2] != y:
244
245
246
247
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
        return pixels
248

249
250
251
    def encode(self, vae, pixels):
        pixels = self.vae_encode_crop_pixels(pixels)
        t = vae.encode(pixels[:,:,:,:3])
252
        return ({"samples":t}, )
comfyanonymous's avatar
comfyanonymous committed
253

comfyanonymous's avatar
comfyanonymous committed
254
255
256
257
258
259
260
261
262
263
class VAEEncodeTiled:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "_for_testing"

    def encode(self, vae, pixels):
264
        pixels = VAEEncode.vae_encode_crop_pixels(pixels)
comfyanonymous's avatar
comfyanonymous committed
265
266
        t = vae.encode_tiled(pixels[:,:,:,:3])
        return ({"samples":t}, )
267

268
269
270
class VAEEncodeForInpaint:
    @classmethod
    def INPUT_TYPES(s):
271
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", ), "mask": ("MASK", ), "grow_mask_by": ("INT", {"default": 6, "min": 0, "max": 64, "step": 1}),}}
272
273
274
275
276
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "latent/inpaint"

277
    def encode(self, vae, pixels, mask, grow_mask_by=6):
278
279
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
280
        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
281

282
        pixels = pixels.clone()
283
        if pixels.shape[1] != x or pixels.shape[2] != y:
284
285
286
287
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
288

289
        #grow mask by a few pixels to keep things seamless in latent space
290
291
292
293
294
295
296
297
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)

298
        m = (1.0 - mask.round()).squeeze(1)
299
300
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
301
            pixels[:,:,:,i] *= m
302
303
304
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

305
        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )
comfyanonymous's avatar
comfyanonymous committed
306

Dr.Lt.Data's avatar
Dr.Lt.Data committed
307
308
class SaveLatent:
    def __init__(self):
309
        self.output_dir = folder_paths.get_output_directory()
Dr.Lt.Data's avatar
Dr.Lt.Data committed
310
311
312
313

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ),
314
                              "filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
Dr.Lt.Data's avatar
Dr.Lt.Data committed
315
316
317
318
319
320
321
322
323
324
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
    RETURN_TYPES = ()
    FUNCTION = "save"

    OUTPUT_NODE = True

    CATEGORY = "_for_testing"

    def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
325
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
Dr.Lt.Data's avatar
Dr.Lt.Data committed
326
327
328
329
330
331

        # support save metadata for latent sharing
        prompt_info = ""
        if prompt is not None:
            prompt_info = json.dumps(prompt)

332
        metadata = {"prompt": prompt_info}
Dr.Lt.Data's avatar
Dr.Lt.Data committed
333
334
335
336
337
338
339
        if extra_pnginfo is not None:
            for x in extra_pnginfo:
                metadata[x] = json.dumps(extra_pnginfo[x])

        file = f"{filename}_{counter:05}_.latent"
        file = os.path.join(full_output_folder, file)

340
341
        output = {}
        output["latent_tensor"] = samples["samples"]
342
        output["latent_format_version_0"] = torch.tensor([])
343

344
        comfy.utils.save_torch_file(output, file, metadata=metadata)
Dr.Lt.Data's avatar
Dr.Lt.Data committed
345
346
347
348
349
350
        return {}


class LoadLatent:
    @classmethod
    def INPUT_TYPES(s):
351
352
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".latent")]
Dr.Lt.Data's avatar
Dr.Lt.Data committed
353
354
355
356
357
358
359
360
        return {"required": {"latent": [sorted(files), ]}, }

    CATEGORY = "_for_testing"

    RETURN_TYPES = ("LATENT", )
    FUNCTION = "load"

    def load(self, latent):
361
362
        latent_path = folder_paths.get_annotated_filepath(latent)
        latent = safetensors.torch.load_file(latent_path, device="cpu")
363
364
365
366
        multiplier = 1.0
        if "latent_format_version_0" not in latent:
            multiplier = 1.0 / 0.18215
        samples = {"samples": latent["latent_tensor"].float() * multiplier}
367
        return (samples, )
Dr.Lt.Data's avatar
Dr.Lt.Data committed
368

369
370
371
372
373
374
375
376
377
378
379
380
381
382
    @classmethod
    def IS_CHANGED(s, latent):
        image_path = folder_paths.get_annotated_filepath(latent)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()

    @classmethod
    def VALIDATE_INPUTS(s, latent):
        if not folder_paths.exists_annotated_filepath(latent):
            return "Invalid latent file: {}".format(latent)
        return True

Dr.Lt.Data's avatar
Dr.Lt.Data committed
383

comfyanonymous's avatar
comfyanonymous committed
384
385
386
class CheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
387
388
        return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ),
                              "ckpt_name": (folder_paths.get_filename_list("checkpoints"), )}}
comfyanonymous's avatar
comfyanonymous committed
389
390
391
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

392
    CATEGORY = "advanced/loaders"
393

comfyanonymous's avatar
comfyanonymous committed
394
    def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
395
396
        config_path = folder_paths.get_full_path("configs", config_name)
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
397
        return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
comfyanonymous's avatar
comfyanonymous committed
398

399
400
401
class CheckpointLoaderSimple:
    @classmethod
    def INPUT_TYPES(s):
402
        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
403
404
405
406
                             }}
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

407
    CATEGORY = "loaders"
408

409
    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
410
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
411
        out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
412
413
        return out

sALTaccount's avatar
sALTaccount committed
414
415
416
class DiffusersLoader:
    @classmethod
    def INPUT_TYPES(cls):
417
        paths = []
sALTaccount's avatar
sALTaccount committed
418
        for search_path in folder_paths.get_folder_paths("diffusers"):
419
            if os.path.exists(search_path):
420
421
422
423
                for root, subdir, files in os.walk(search_path, followlinks=True):
                    if "model_index.json" in files:
                        paths.append(os.path.relpath(root, start=search_path))

424
        return {"required": {"model_path": (paths,), }}
sALTaccount's avatar
sALTaccount committed
425
426
427
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

428
    CATEGORY = "advanced/loaders/deprecated"
sALTaccount's avatar
sALTaccount committed
429
430

    def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
sALTaccount's avatar
sALTaccount committed
431
432
        for search_path in folder_paths.get_folder_paths("diffusers"):
            if os.path.exists(search_path):
433
434
435
                path = os.path.join(search_path, model_path)
                if os.path.exists(path):
                    model_path = path
sALTaccount's avatar
sALTaccount committed
436
                    break
437

438
        return comfy.diffusers_load.load_diffusers(model_path, fp16=comfy.model_management.should_use_fp16(), output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
sALTaccount's avatar
sALTaccount committed
439
440


441
442
443
444
445
446
447
448
class unCLIPCheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
                             }}
    RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION")
    FUNCTION = "load_checkpoint"

449
    CATEGORY = "loaders"
450
451
452
453
454
455

    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
        out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
        return out

comfyanonymous's avatar
comfyanonymous committed
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
class CLIPSetLastLayer:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip": ("CLIP", ),
                              "stop_at_clip_layer": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
                              }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "set_last_layer"

    CATEGORY = "conditioning"

    def set_last_layer(self, clip, stop_at_clip_layer):
        clip = clip.clone()
        clip.clip_layer(stop_at_clip_layer)
        return (clip,)

472
class LoraLoader:
473
474
475
    def __init__(self):
        self.loaded_lora = None

476
477
478
479
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "clip": ("CLIP", ),
480
                              "lora_name": (folder_paths.get_filename_list("loras"), ),
481
482
                              "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
483
484
485
486
487
488
489
                              }}
    RETURN_TYPES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"

    CATEGORY = "loaders"

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
490
491
492
        if strength_model == 0 and strength_clip == 0:
            return (model, clip)

493
        lora_path = folder_paths.get_full_path("loras", lora_name)
494
495
496
497
498
499
500
501
502
503
504
505
        lora = None
        if self.loaded_lora is not None:
            if self.loaded_lora[0] == lora_path:
                lora = self.loaded_lora[1]
            else:
                del self.loaded_lora

        if lora is None:
            lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
            self.loaded_lora = (lora_path, lora)

        model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
506
507
        return (model_lora, clip_lora)

comfyanonymous's avatar
comfyanonymous committed
508
509
510
class VAELoader:
    @classmethod
    def INPUT_TYPES(s):
511
        return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}}
comfyanonymous's avatar
comfyanonymous committed
512
513
514
    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"

515
516
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
517
518
    #TODO: scale factor?
    def load_vae(self, vae_name):
519
        vae_path = folder_paths.get_full_path("vae", vae_name)
comfyanonymous's avatar
comfyanonymous committed
520
521
522
        vae = comfy.sd.VAE(ckpt_path=vae_path)
        return (vae,)

comfyanonymous's avatar
comfyanonymous committed
523
524
525
class ControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
526
        return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
comfyanonymous's avatar
comfyanonymous committed
527
528
529
530
531
532
533

    RETURN_TYPES = ("CONTROL_NET",)
    FUNCTION = "load_controlnet"

    CATEGORY = "loaders"

    def load_controlnet(self, control_net_name):
534
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
comfyanonymous's avatar
comfyanonymous committed
535
536
537
        controlnet = comfy.sd.load_controlnet(controlnet_path)
        return (controlnet,)

538
539
540
541
class DiffControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
542
                              "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
543
544
545
546
547
548
549

    RETURN_TYPES = ("CONTROL_NET",)
    FUNCTION = "load_controlnet"

    CATEGORY = "loaders"

    def load_controlnet(self, model, control_net_name):
550
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
551
552
553
        controlnet = comfy.sd.load_controlnet(controlnet_path, model)
        return (controlnet,)

comfyanonymous's avatar
comfyanonymous committed
554
555
556
557

class ControlNetApply:
    @classmethod
    def INPUT_TYPES(s):
558
559
560
561
562
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "control_net": ("CONTROL_NET", ),
                             "image": ("IMAGE", ),
                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01})
                             }}
comfyanonymous's avatar
comfyanonymous committed
563
564
565
566
567
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

568
    def apply_controlnet(self, conditioning, control_net, image, strength):
569
570
571
        if strength == 0:
            return (conditioning, )

comfyanonymous's avatar
comfyanonymous committed
572
573
574
575
        c = []
        control_hint = image.movedim(-1,1)
        for t in conditioning:
            n = [t[0], t[1].copy()]
comfyanonymous's avatar
comfyanonymous committed
576
577
578
579
            c_net = control_net.copy().set_cond_hint(control_hint, strength)
            if 'control' in t[1]:
                c_net.set_previous_controlnet(t[1]['control'])
            n[1]['control'] = c_net
comfyanonymous's avatar
comfyanonymous committed
580
581
582
            c.append(n)
        return (c, )

583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
class UNETLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ),
                             }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"

    CATEGORY = "advanced/loaders"

    def load_unet(self, unet_name):
        unet_path = folder_paths.get_full_path("unet", unet_name)
        model = comfy.sd.load_unet(unet_path)
        return (model,)

598
599
600
class CLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
601
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
602
603
604
605
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

606
    CATEGORY = "advanced/loaders"
607

608
    def load_clip(self, clip_name):
609
        clip_path = folder_paths.get_full_path("clip", clip_name)
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
        clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"))
        return (clip,)

class DualCLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip_name1": (folder_paths.get_filename_list("clip"), ), "clip_name2": (folder_paths.get_filename_list("clip"), ),
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

    CATEGORY = "advanced/loaders"

    def load_clip(self, clip_name1, clip_name2):
        clip_path1 = folder_paths.get_full_path("clip", clip_name1)
        clip_path2 = folder_paths.get_full_path("clip", clip_name2)
        clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=folder_paths.get_folder_paths("embeddings"))
627
628
        return (clip,)

629
630
631
class CLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
632
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
633
634
635
636
637
638
639
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

    def load_clip(self, clip_name):
640
        clip_path = folder_paths.get_full_path("clip_vision", clip_name)
641
        clip_vision = comfy.clip_vision.load(clip_path)
642
643
644
645
646
647
648
649
        return (clip_vision,)

class CLIPVisionEncode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip_vision": ("CLIP_VISION",),
                              "image": ("IMAGE",)
                             }}
comfyanonymous's avatar
comfyanonymous committed
650
    RETURN_TYPES = ("CLIP_VISION_OUTPUT",)
651
652
    FUNCTION = "encode"

653
    CATEGORY = "conditioning"
654
655
656
657
658
659
660
661

    def encode(self, clip_vision, image):
        output = clip_vision.encode_image(image)
        return (output,)

class StyleModelLoader:
    @classmethod
    def INPUT_TYPES(s):
662
        return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}}
663
664
665
666
667
668
669

    RETURN_TYPES = ("STYLE_MODEL",)
    FUNCTION = "load_style_model"

    CATEGORY = "loaders"

    def load_style_model(self, style_model_name):
670
        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
671
672
673
674
675
676
677
        style_model = comfy.sd.load_style_model(style_model_path)
        return (style_model,)


class StyleModelApply:
    @classmethod
    def INPUT_TYPES(s):
678
679
680
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "style_model": ("STYLE_MODEL", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
681
682
683
684
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_stylemodel"

comfyanonymous's avatar
comfyanonymous committed
685
    CATEGORY = "conditioning/style_model"
686

687
688
    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
        cond = style_model.get_cond(clip_vision_output)
689
        c = []
690
691
        for t in conditioning:
            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
692
693
694
            c.append(n)
        return (c, )

695
696
697
698
699
700
class unCLIPConditioning:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
                             "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
701
                             "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
702
703
704
705
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_adm"

706
    CATEGORY = "conditioning"
707

708
    def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation):
709
710
711
        if strength == 0:
            return (conditioning, )

712
713
714
        c = []
        for t in conditioning:
            o = t[1].copy()
715
716
717
            x = {"clip_vision_output": clip_vision_output, "strength": strength, "noise_augmentation": noise_augmentation}
            if "unclip_conditioning" in o:
                o["unclip_conditioning"] = o["unclip_conditioning"][:] + [x]
718
            else:
719
                o["unclip_conditioning"] = [x]
720
721
722
723
            n = [t[0], o]
            c.append(n)
        return (c, )

724
725
726
727
728
729
730
731
class GLIGENLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "gligen_name": (folder_paths.get_filename_list("gligen"), )}}

    RETURN_TYPES = ("GLIGEN",)
    FUNCTION = "load_gligen"

comfyanonymous's avatar
comfyanonymous committed
732
    CATEGORY = "loaders"
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753

    def load_gligen(self, gligen_name):
        gligen_path = folder_paths.get_full_path("gligen", gligen_name)
        gligen = comfy.sd.load_gligen(gligen_path)
        return (gligen,)

class GLIGENTextBoxApply:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning_to": ("CONDITIONING", ),
                              "clip": ("CLIP", ),
                              "gligen_textbox_model": ("GLIGEN", ),
                              "text": ("STRING", {"multiline": True}),
                              "width": ("INT", {"default": 64, "min": 8, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 64, "min": 8, "max": MAX_RESOLUTION, "step": 8}),
                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "append"

comfyanonymous's avatar
comfyanonymous committed
754
    CATEGORY = "conditioning/gligen"
755
756
757
758
759
760
761
762
763
764
765
766
767
768

    def append(self, conditioning_to, clip, gligen_textbox_model, text, width, height, x, y):
        c = []
        cond, cond_pooled = clip.encode_from_tokens(clip.tokenize(text), return_pooled=True)
        for t in conditioning_to:
            n = [t[0], t[1].copy()]
            position_params = [(cond_pooled, height // 8, width // 8, y // 8, x // 8)]
            prev = []
            if "gligen" in n[1]:
                prev = n[1]['gligen'][2]

            n[1]['gligen'] = ("position", gligen_textbox_model, prev + position_params)
            c.append(n)
        return (c, )
769

comfyanonymous's avatar
comfyanonymous committed
770
771
772
773
774
775
class EmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
776
777
        return {"required": { "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
comfyanonymous's avatar
comfyanonymous committed
778
779
780
781
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

782
783
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
784
785
    def generate(self, width, height, batch_size=1):
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
786
        return ({"samples":latent}, )
comfyanonymous's avatar
comfyanonymous committed
787

comfyanonymous's avatar
comfyanonymous committed
788

789
790
791
792
793
class LatentFromBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "batch_index": ("INT", {"default": 0, "min": 0, "max": 63}),
794
                              "length": ("INT", {"default": 1, "min": 1, "max": 64}),
795
796
                              }}
    RETURN_TYPES = ("LATENT",)
797
    FUNCTION = "frombatch"
798

799
    CATEGORY = "latent/batch"
800

801
    def frombatch(self, samples, batch_index, length):
802
803
804
        s = samples.copy()
        s_in = samples["samples"]
        batch_index = min(s_in.shape[0] - 1, batch_index)
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
        length = min(s_in.shape[0] - batch_index, length)
        s["samples"] = s_in[batch_index:batch_index + length].clone()
        if "noise_mask" in samples:
            masks = samples["noise_mask"]
            if masks.shape[0] == 1:
                s["noise_mask"] = masks.clone()
            else:
                if masks.shape[0] < s_in.shape[0]:
                    masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
                s["noise_mask"] = masks[batch_index:batch_index + length].clone()
        if "batch_index" not in s:
            s["batch_index"] = [x for x in range(batch_index, batch_index+length)]
        else:
            s["batch_index"] = samples["batch_index"][batch_index:batch_index + length]
        return (s,)
    
class RepeatLatentBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "amount": ("INT", {"default": 1, "min": 1, "max": 64}),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "repeat"

    CATEGORY = "latent/batch"

    def repeat(self, samples, amount):
        s = samples.copy()
        s_in = samples["samples"]
        
        s["samples"] = s_in.repeat((amount, 1,1,1))
        if "noise_mask" in samples and samples["noise_mask"].shape[0] > 1:
            masks = samples["noise_mask"]
            if masks.shape[0] < s_in.shape[0]:
                masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
            s["noise_mask"] = samples["noise_mask"].repeat((amount, 1,1,1))
        if "batch_index" in s:
            offset = max(s["batch_index"]) - min(s["batch_index"]) + 1
            s["batch_index"] = s["batch_index"] + [x + (i * offset) for i in range(1, amount) for x in s["batch_index"]]
845
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
846

comfyanonymous's avatar
comfyanonymous committed
847
class LatentUpscale:
comfyanonymous's avatar
comfyanonymous committed
848
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
849
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
850
851
852
853

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
854
855
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
856
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
857
858
859
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

860
861
    CATEGORY = "latent"

862
    def upscale(self, samples, upscale_method, width, height, crop):
863
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
864
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
865
866
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
867
class LatentUpscaleBy:
comfyanonymous's avatar
comfyanonymous committed
868
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
comfyanonymous's avatar
comfyanonymous committed
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
                              "scale_by": ("FLOAT", {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}),}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

    CATEGORY = "latent"

    def upscale(self, samples, upscale_method, scale_by):
        s = samples.copy()
        width = round(samples["samples"].shape[3] * scale_by)
        height = round(samples["samples"].shape[2] * scale_by)
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width, height, upscale_method, "disabled")
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
886
887
888
889
890
891
892
893
894
class LatentRotate:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "rotation": (["none", "90 degrees", "180 degrees", "270 degrees"],),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "rotate"

comfyanonymous's avatar
comfyanonymous committed
895
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
896
897

    def rotate(self, samples, rotation):
898
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
899
900
901
902
903
904
905
906
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

907
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
908
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
909
910
911
912
913
914
915
916
917
918

class LatentFlip:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "flip_method": (["x-axis: vertically", "y-axis: horizontally"],),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "flip"

comfyanonymous's avatar
comfyanonymous committed
919
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
920
921

    def flip(self, samples, flip_method):
922
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
923
        if flip_method.startswith("x"):
924
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
925
        elif flip_method.startswith("y"):
926
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
927
928

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
929
930
931
932

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
933
934
935
936
937
938
        return {"required": { "samples_to": ("LATENT",),
                              "samples_from": ("LATENT",),
                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              "feather": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              }}
comfyanonymous's avatar
comfyanonymous committed
939
940
941
942
943
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
944
945
946
    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
        x =  x // 8
        y = y // 8
947
        feather = feather // 8
comfyanonymous's avatar
comfyanonymous committed
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
        samples_out = samples_to.copy()
        s = samples_to["samples"].clone()
        samples_to = samples_to["samples"]
        samples_from = samples_from["samples"]
        if feather == 0:
            s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
        else:
            samples_from = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
            mask = torch.ones_like(samples_from)
            for t in range(feather):
                if y != 0:
                    mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))

                if y + samples_from.shape[2] < samples_to.shape[2]:
                    mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
                if x != 0:
                    mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
                if x + samples_from.shape[3] < samples_to.shape[3]:
                    mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
            rev_mask = torch.ones_like(mask) - mask
            s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x] * mask + s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] * rev_mask
        samples_out["samples"] = s
        return (samples_out,)
comfyanonymous's avatar
comfyanonymous committed
971

comfyanonymous's avatar
comfyanonymous committed
972
973
974
975
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
976
977
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
978
979
                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
comfyanonymous's avatar
comfyanonymous committed
980
981
982
983
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
984
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
985
986

    def crop(self, samples, width, height, x, y):
987
988
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
989
990
991
992
993
994
995
996
997
998
999
1000
1001
        x =  x // 8
        y = y // 8

        #enfonce minimum size of 64
        if x > (samples.shape[3] - 8):
            x = samples.shape[3] - 8
        if y > (samples.shape[2] - 8):
            y = samples.shape[2] - 8

        new_height = height // 8
        new_width = width // 8
        to_x = new_width + x
        to_y = new_height + y
1002
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
1003
1004
        return (s,)

1005
1006
1007
1008
1009
1010
1011
1012
1013
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

1014
    CATEGORY = "latent/inpaint"
1015
1016
1017

    def set_mask(self, samples, mask):
        s = samples.copy()
1018
        s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
1019
1020
        return (s,)

space-nuko's avatar
space-nuko committed
1021

space-nuko's avatar
space-nuko committed
1022
def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
1023
    device = comfy.model_management.get_torch_device()
1024
    latent_image = latent["samples"]
1025

comfyanonymous's avatar
comfyanonymous committed
1026
1027
1028
    if disable_noise:
        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
    else:
1029
1030
        batch_inds = latent["batch_index"] if "batch_index" in latent else None
        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
comfyanonymous's avatar
comfyanonymous committed
1031

1032
    noise_mask = None
1033
    if "noise_mask" in latent:
1034
        noise_mask = latent["noise_mask"]
comfyanonymous's avatar
comfyanonymous committed
1035

space-nuko's avatar
space-nuko committed
1036
1037
1038
1039
    preview_format = "JPEG"
    if preview_format not in ["JPEG", "PNG"]:
        preview_format = "JPEG"

1040
    previewer = latent_preview.get_previewer(device, model.model.latent_format)
space-nuko's avatar
space-nuko committed
1041

1042
    pbar = comfy.utils.ProgressBar(steps)
1043
    def callback(step, x0, x, total_steps):
space-nuko's avatar
space-nuko committed
1044
        preview_bytes = None
1045
        if previewer:
1046
            preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
space-nuko's avatar
space-nuko committed
1047
        pbar.update_absolute(step + 1, total_steps, preview_bytes)
1048

1049
1050
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
1051
                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, seed=seed)
1052
1053
1054
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
1055

comfyanonymous's avatar
comfyanonymous committed
1056
1057
1058
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
1059
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
                    {"model": ("MODEL",),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "latent_image": ("LATENT", ),
                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
space-nuko's avatar
space-nuko committed
1070
1071
                     }
                }
comfyanonymous's avatar
comfyanonymous committed
1072
1073
1074
1075

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "sample"

1076
1077
    CATEGORY = "sampling"

space-nuko's avatar
space-nuko committed
1078
1079
    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
comfyanonymous's avatar
comfyanonymous committed
1080

comfyanonymous's avatar
comfyanonymous committed
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
class KSamplerAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL",),
                    "add_noise": (["enable", "disable"], ),
                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "latent_image": ("LATENT", ),
                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
                    "return_with_leftover_noise": (["disable", "enable"], ),
space-nuko's avatar
space-nuko committed
1098
1099
                     }
                }
comfyanonymous's avatar
comfyanonymous committed
1100
1101
1102
1103
1104

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "sample"

    CATEGORY = "sampling"
comfyanonymous's avatar
comfyanonymous committed
1105

space-nuko's avatar
space-nuko committed
1106
    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
comfyanonymous's avatar
comfyanonymous committed
1107
1108
1109
1110
1111
1112
        force_full_denoise = True
        if return_with_leftover_noise == "enable":
            force_full_denoise = False
        disable_noise = False
        if add_noise == "disable":
            disable_noise = True
space-nuko's avatar
space-nuko committed
1113
        return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
comfyanonymous's avatar
comfyanonymous committed
1114
1115
1116

class SaveImage:
    def __init__(self):
1117
        self.output_dir = folder_paths.get_output_directory()
m957ymj75urz's avatar
m957ymj75urz committed
1118
        self.type = "output"
comfyanonymous's avatar
comfyanonymous committed
1119
1120
1121
1122

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
1123
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
1124
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
1125
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
1126
1127
1128
1129
1130
1131
1132
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

1133
1134
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
1135
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
1136
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
m957ymj75urz's avatar
m957ymj75urz committed
1137
        results = list()
comfyanonymous's avatar
comfyanonymous committed
1138
1139
        for image in images:
            i = 255. * image.cpu().numpy()
1140
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
comfyanonymous's avatar
comfyanonymous committed
1141
1142
1143
1144
1145
1146
            metadata = PngInfo()
            if prompt is not None:
                metadata.add_text("prompt", json.dumps(prompt))
            if extra_pnginfo is not None:
                for x in extra_pnginfo:
                    metadata.add_text(x, json.dumps(extra_pnginfo[x]))
1147

1148
            file = f"{filename}_{counter:05}_.png"
1149
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
m957ymj75urz's avatar
m957ymj75urz committed
1150
1151
1152
1153
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
Gavroche CryptoRUSH's avatar
Gavroche CryptoRUSH committed
1154
            })
1155
            counter += 1
comfyanonymous's avatar
comfyanonymous committed
1156

m957ymj75urz's avatar
m957ymj75urz committed
1157
        return { "ui": { "images": results } }
comfyanonymous's avatar
comfyanonymous committed
1158

pythongosssss's avatar
pythongosssss committed
1159
1160
class PreviewImage(SaveImage):
    def __init__(self):
1161
        self.output_dir = folder_paths.get_temp_directory()
m957ymj75urz's avatar
m957ymj75urz committed
1162
        self.type = "temp"
pythongosssss's avatar
pythongosssss committed
1163
1164
1165

    @classmethod
    def INPUT_TYPES(s):
1166
        return {"required":
pythongosssss's avatar
pythongosssss committed
1167
1168
1169
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
1170

1171
1172
1173
class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
1174
        input_dir = folder_paths.get_input_directory()
comfyanonymous's avatar
comfyanonymous committed
1175
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
1176
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
1177
                    {"image": (sorted(files), )},
1178
                }
1179
1180

    CATEGORY = "image"
1181

1182
    RETURN_TYPES = ("IMAGE", "MASK")
1183
1184
    FUNCTION = "load_image"
    def load_image(self, image):
1185
        image_path = folder_paths.get_annotated_filepath(image)
1186
        i = Image.open(image_path)
1187
        i = ImageOps.exif_transpose(i)
1188
        image = i.convert("RGB")
1189
        image = np.array(image).astype(np.float32) / 255.0
1190
        image = torch.from_numpy(image)[None,]
1191
1192
1193
1194
1195
1196
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
        return (image, mask)
1197

1198
1199
    @classmethod
    def IS_CHANGED(s, image):
1200
        image_path = folder_paths.get_annotated_filepath(image)
1201
1202
1203
1204
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
1205

1206
1207
1208
1209
1210
1211
1212
    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        return True

1213
class LoadImageMask:
1214
    _color_channels = ["alpha", "red", "green", "blue"]
1215
1216
    @classmethod
    def INPUT_TYPES(s):
1217
        input_dir = folder_paths.get_input_directory()
comfyanonymous's avatar
comfyanonymous committed
1218
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
1219
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
1220
                    {"image": (sorted(files), ),
1221
                     "channel": (s._color_channels, ), }
1222
1223
                }

1224
    CATEGORY = "mask"
1225
1226
1227
1228

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
1229
        image_path = folder_paths.get_annotated_filepath(image)
1230
        i = Image.open(image_path)
1231
        i = ImageOps.exif_transpose(i)
1232
1233
        if i.getbands() != ("R", "G", "B", "A"):
            i = i.convert("RGBA")
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
        mask = None
        c = channel[0].upper()
        if c in i.getbands():
            mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
            mask = torch.from_numpy(mask)
            if c == 'A':
                mask = 1. - mask
        else:
            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
        return (mask,)

    @classmethod
    def IS_CHANGED(s, image, channel):
1247
        image_path = folder_paths.get_annotated_filepath(image)
1248
1249
1250
1251
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
1252

1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
    @classmethod
    def VALIDATE_INPUTS(s, image, channel):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        if channel not in s._color_channels:
            return "Invalid color channel: {}".format(channel)

        return True

comfyanonymous's avatar
comfyanonymous committed
1263
class ImageScale:
comfyanonymous's avatar
comfyanonymous committed
1264
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic"]
comfyanonymous's avatar
comfyanonymous committed
1265
1266
1267
1268
1269
    crop_methods = ["disabled", "center"]

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
1270
1271
                              "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
                              "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
comfyanonymous's avatar
comfyanonymous committed
1272
1273
1274
1275
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

1276
    CATEGORY = "image/upscaling"
1277

comfyanonymous's avatar
comfyanonymous committed
1278
1279
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
1280
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
1281
1282
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
1283

comfyanonymous's avatar
comfyanonymous committed
1284
class ImageScaleBy:
comfyanonymous's avatar
comfyanonymous committed
1285
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic"]
comfyanonymous's avatar
comfyanonymous committed
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
                              "scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}),}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

    CATEGORY = "image/upscaling"

    def upscale(self, image, upscale_method, scale_by):
        samples = image.movedim(-1,1)
        width = round(samples.shape[3] * scale_by)
        height = round(samples.shape[2] * scale_by)
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
        s = s.movedim(1,-1)
        return (s,)

1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
class ImageInvert:

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",)}}

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "invert"

    CATEGORY = "image"

    def invert(self, image):
        s = 1.0 - image
        return (s,)


Guo Y.K's avatar
Guo Y.K committed
1320
1321
1322
1323
1324
1325
1326
class ImagePadForOutpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
1327
1328
1329
1330
                "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
1331
                "feathering": ("INT", {"default": 40, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
Guo Y.K's avatar
Guo Y.K committed
1332
1333
1334
1335
1336
1337
1338
1339
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "expand_image"

    CATEGORY = "image"

1340
    def expand_image(self, image, left, top, right, bottom, feathering):
Guo Y.K's avatar
Guo Y.K committed
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
        d1, d2, d3, d4 = image.size()

        new_image = torch.zeros(
            (d1, d2 + top + bottom, d3 + left + right, d4),
            dtype=torch.float32,
        )
        new_image[:, top:top + d2, left:left + d3, :] = image

        mask = torch.ones(
            (d2 + top + bottom, d3 + left + right),
            dtype=torch.float32,
        )
1353

1354
1355
1356
1357
1358
        t = torch.zeros(
            (d2, d3),
            dtype=torch.float32
        )

1359
        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378

            for i in range(d2):
                for j in range(d3):
                    dt = i if top != 0 else d2
                    db = d2 - i if bottom != 0 else d2

                    dl = j if left != 0 else d3
                    dr = d3 - j if right != 0 else d3

                    d = min(dt, db, dl, dr)

                    if d >= feathering:
                        continue

                    v = (feathering - d) / feathering

                    t[i, j] = v * v

        mask[top:top + d2, left:left + d3] = t
1379

Guo Y.K's avatar
Guo Y.K committed
1380
1381
1382
        return (new_image, mask)


comfyanonymous's avatar
comfyanonymous committed
1383
1384
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
1385
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
1386
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
1387
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
1388
1389
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
1390
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
1391
1392
1393
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
    "LatentUpscale": LatentUpscale,
comfyanonymous's avatar
comfyanonymous committed
1394
    "LatentUpscaleBy": LatentUpscaleBy,
1395
    "LatentFromBatch": LatentFromBatch,
1396
    "RepeatLatentBatch": RepeatLatentBatch,
comfyanonymous's avatar
comfyanonymous committed
1397
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
1398
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
1399
    "LoadImage": LoadImage,
1400
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
1401
    "ImageScale": ImageScale,
comfyanonymous's avatar
comfyanonymous committed
1402
    "ImageScaleBy": ImageScaleBy,
1403
    "ImageInvert": ImageInvert,
Guo Y.K's avatar
Guo Y.K committed
1404
    "ImagePadForOutpaint": ImagePadForOutpaint,
FizzleDorf's avatar
FizzleDorf committed
1405
    "ConditioningAverage ": ConditioningAverage ,
comfyanonymous's avatar
comfyanonymous committed
1406
1407
    "ConditioningCombine": ConditioningCombine,
    "ConditioningSetArea": ConditioningSetArea,
Jacob Segal's avatar
Jacob Segal committed
1408
    "ConditioningSetMask": ConditioningSetMask,
comfyanonymous's avatar
comfyanonymous committed
1409
    "KSamplerAdvanced": KSamplerAdvanced,
1410
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
1411
    "LatentComposite": LatentComposite,
comfyanonymous's avatar
comfyanonymous committed
1412
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
1413
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
1414
    "LatentCrop": LatentCrop,
1415
    "LoraLoader": LoraLoader,
1416
    "CLIPLoader": CLIPLoader,
1417
    "UNETLoader": UNETLoader,
1418
    "DualCLIPLoader": DualCLIPLoader,
1419
    "CLIPVisionEncode": CLIPVisionEncode,
1420
    "StyleModelApply": StyleModelApply,
1421
    "unCLIPConditioning": unCLIPConditioning,
comfyanonymous's avatar
comfyanonymous committed
1422
1423
    "ControlNetApply": ControlNetApply,
    "ControlNetLoader": ControlNetLoader,
1424
    "DiffControlNetLoader": DiffControlNetLoader,
comfyanonymous's avatar
comfyanonymous committed
1425
1426
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
1427
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
1428
    "VAEEncodeTiled": VAEEncodeTiled,
1429
    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
1430
1431
1432
    "GLIGENLoader": GLIGENLoader,
    "GLIGENTextBoxApply": GLIGENTextBoxApply,

1433
    "CheckpointLoader": CheckpointLoader,
sALTaccount's avatar
sALTaccount committed
1434
    "DiffusersLoader": DiffusersLoader,
Dr.Lt.Data's avatar
Dr.Lt.Data committed
1435
1436

    "LoadLatent": LoadLatent,
1437
    "SaveLatent": SaveLatent,
1438
1439

    "ConditioningZeroOut": ConditioningZeroOut,
1440
    "ConditioningConcat": ConditioningConcat,
comfyanonymous's avatar
comfyanonymous committed
1441
1442
}

City's avatar
City committed
1443
1444
1445
1446
1447
NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "KSampler": "KSampler",
    "KSamplerAdvanced": "KSampler (Advanced)",
    # Loaders
comfyanonymous's avatar
Rename.  
comfyanonymous committed
1448
1449
    "CheckpointLoader": "Load Checkpoint (With Config)",
    "CheckpointLoaderSimple": "Load Checkpoint",
City's avatar
City committed
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
    "VAELoader": "Load VAE",
    "LoraLoader": "Load LoRA",
    "CLIPLoader": "Load CLIP",
    "ControlNetLoader": "Load ControlNet Model",
    "DiffControlNetLoader": "Load ControlNet Model (diff)",
    "StyleModelLoader": "Load Style Model",
    "CLIPVisionLoader": "Load CLIP Vision",
    "UpscaleModelLoader": "Load Upscale Model",
    # Conditioning
    "CLIPVisionEncode": "CLIP Vision Encode",
    "StyleModelApply": "Apply Style Model",
    "CLIPTextEncode": "CLIP Text Encode (Prompt)",
    "CLIPSetLastLayer": "CLIP Set Last Layer",
    "ConditioningCombine": "Conditioning (Combine)",
FizzleDorf's avatar
FizzleDorf committed
1464
    "ConditioningAverage ": "Conditioning (Average)",
City's avatar
City committed
1465
    "ConditioningSetArea": "Conditioning (Set Area)",
Jacob Segal's avatar
Jacob Segal committed
1466
    "ConditioningSetMask": "Conditioning (Set Mask)",
City's avatar
City committed
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
    "ControlNetApply": "Apply ControlNet",
    # Latent
    "VAEEncodeForInpaint": "VAE Encode (for Inpainting)",
    "SetLatentNoiseMask": "Set Latent Noise Mask",
    "VAEDecode": "VAE Decode",
    "VAEEncode": "VAE Encode",
    "LatentRotate": "Rotate Latent",
    "LatentFlip": "Flip Latent",
    "LatentCrop": "Crop Latent",
    "EmptyLatentImage": "Empty Latent Image",
    "LatentUpscale": "Upscale Latent",
comfyanonymous's avatar
comfyanonymous committed
1478
    "LatentUpscaleBy": "Upscale Latent By",
City's avatar
City committed
1479
    "LatentComposite": "Latent Composite",
1480
1481
    "LatentFromBatch" : "Latent From Batch",
    "RepeatLatentBatch": "Repeat Latent Batch",
City's avatar
City committed
1482
1483
1484
1485
1486
1487
    # Image
    "SaveImage": "Save Image",
    "PreviewImage": "Preview Image",
    "LoadImage": "Load Image",
    "LoadImageMask": "Load Image (as Mask)",
    "ImageScale": "Upscale Image",
comfyanonymous's avatar
comfyanonymous committed
1488
    "ImageScaleBy": "Upscale Image By",
City's avatar
City committed
1489
1490
1491
1492
1493
1494
1495
1496
    "ImageUpscaleWithModel": "Upscale Image (using Model)",
    "ImageInvert": "Invert Image",
    "ImagePadForOutpaint": "Pad Image for Outpainting",
    # _for_testing
    "VAEDecodeTiled": "VAE Decode (Tiled)",
    "VAEEncodeTiled": "VAE Encode (Tiled)",
}

1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
def load_custom_node(module_path):
    module_name = os.path.basename(module_path)
    if os.path.isfile(module_path):
        sp = os.path.splitext(module_path)
        module_name = sp[0]
    try:
        if os.path.isfile(module_path):
            module_spec = importlib.util.spec_from_file_location(module_name, module_path)
        else:
            module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
        module = importlib.util.module_from_spec(module_spec)
        sys.modules[module_name] = module
        module_spec.loader.exec_module(module)
        if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
            NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
1512
1513
            if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
                NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
1514
            return True
1515
1516
        else:
            print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
1517
            return False
1518
1519
1520
    except Exception as e:
        print(traceback.format_exc())
        print(f"Cannot import {module_path} module for custom nodes:", e)
1521
        return False
1522

Hacker 17082006's avatar
Hacker 17082006 committed
1523
def load_custom_nodes():
1524
    node_paths = folder_paths.get_folder_paths("custom_nodes")
1525
    node_import_times = []
1526
1527
1528
1529
1530
1531
1532
1533
    for custom_node_path in node_paths:
        possible_modules = os.listdir(custom_node_path)
        if "__pycache__" in possible_modules:
            possible_modules.remove("__pycache__")

        for possible_module in possible_modules:
            module_path = os.path.join(custom_node_path, possible_module)
            if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
1534
            if module_path.endswith(".disabled"): continue
1535
            time_before = time.perf_counter()
1536
            success = load_custom_node(module_path)
1537
            node_import_times.append((time.perf_counter() - time_before, module_path, success))
1538

1539
    if len(node_import_times) > 0:
comfyanonymous's avatar
comfyanonymous committed
1540
        print("\nImport times for custom nodes:")
1541
        for n in sorted(node_import_times):
1542
1543
1544
1545
1546
            if n[2]:
                import_message = ""
            else:
                import_message = " (IMPORT FAILED)"
            print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
1547
        print()
1548

1549
def init_custom_nodes():
1550
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_hypernetwork.py"))
1551
1552
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py"))
1553
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))
1554
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_rebatch.py"))
1555
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_model_merging.py"))
1556
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_tomesd.py"))
1557
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_clip_sdxl.py"))
1558
    load_custom_nodes()