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

import os
import sys
import json
6
import hashlib
comfyanonymous's avatar
comfyanonymous committed
7
import copy
8
import traceback
comfyanonymous's avatar
comfyanonymous committed
9
10
11
12
13

from PIL import Image
from PIL.PngImagePlugin import PngInfo
import numpy as np

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


import comfy.samplers
import comfy.sd
comfyanonymous's avatar
comfyanonymous committed
19
20
import comfy.utils

21
22
import comfy_extras.clip_vision

23
import model_management
24
import importlib
comfyanonymous's avatar
comfyanonymous committed
25

26
import folder_paths
comfyanonymous's avatar
comfyanonymous committed
27
28
supported_ckpt_extensions = ['.ckpt', '.pth']
supported_pt_extensions = ['.ckpt', '.pt', '.bin', '.pth']
comfyanonymous's avatar
comfyanonymous committed
29
30
31
try:
    import safetensors.torch
    supported_ckpt_extensions += ['.safetensors']
comfyanonymous's avatar
comfyanonymous committed
32
    supported_pt_extensions += ['.safetensors']
comfyanonymous's avatar
comfyanonymous committed
33
34
35
except:
    print("Could not import safetensors, safetensors support disabled.")

36
37
38
39
def recursive_search(directory):  
    result = []
    for root, subdir, file in os.walk(directory, followlinks=True):
        for filepath in file:
40
41
            #we os.path,join directory with a blank string to generate a path separator at the end.
            result.append(os.path.join(root, filepath).replace(os.path.join(directory,''),'')) 
42
43
    return result

comfyanonymous's avatar
comfyanonymous committed
44
45
46
def filter_files_extensions(files, extensions):
    return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files)))

47
48
49
50

def before_node_execution():
    model_management.throw_exception_if_processing_interrupted()

51
52
def interrupt_processing(value=True):
    model_management.interrupt_current_processing(value)
53

comfyanonymous's avatar
comfyanonymous committed
54
55
56
class CLIPTextEncode:
    @classmethod
    def INPUT_TYPES(s):
57
        return {"required": {"text": ("STRING", {"multiline": True}), "clip": ("CLIP", )}}
comfyanonymous's avatar
comfyanonymous committed
58
59
60
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"

61
62
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
63
    def encode(self, clip, text):
comfyanonymous's avatar
comfyanonymous committed
64
65
66
67
68
69
70
71
72
        return ([[clip.encode(text), {}]], )

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

73
74
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
    def combine(self, conditioning_1, conditioning_2):
        return (conditioning_1 + conditioning_2, )

class ConditioningSetArea:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
                              "width": ("INT", {"default": 64, "min": 64, "max": 4096, "step": 64}),
                              "height": ("INT", {"default": 64, "min": 64, "max": 4096, "step": 64}),
                              "x": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
                              "y": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "append"

91
92
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
93
    def append(self, conditioning, width, height, x, y, strength, min_sigma=0.0, max_sigma=99.0):
comfyanonymous's avatar
comfyanonymous committed
94
95
96
97
98
99
100
101
        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
            n[1]['min_sigma'] = min_sigma
            n[1]['max_sigma'] = max_sigma
            c.append(n)
comfyanonymous's avatar
comfyanonymous committed
102
        return (c, )
comfyanonymous's avatar
comfyanonymous committed
103
104
105
106
107
108
109
110
111
112
113

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

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"

114
115
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
116
    def decode(self, vae, samples):
117
        return (vae.decode(samples["samples"]), )
comfyanonymous's avatar
comfyanonymous committed
118

119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class VAEDecodeTiled:
    def __init__(self, device="cpu"):
        self.device = device

    @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
134
135
136
137
138
139
140
141
142
143
class VAEEncode:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

144
145
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
146
    def encode(self, vae, pixels):
147
148
149
150
        x = (pixels.shape[1] // 64) * 64
        y = (pixels.shape[2] // 64) * 64
        if pixels.shape[1] != x or pixels.shape[2] != y:
            pixels = pixels[:,:x,:y,:]
151
152
153
        t = vae.encode(pixels[:,:,:,:3])

        return ({"samples":t}, )
comfyanonymous's avatar
comfyanonymous committed
154

comfyanonymous's avatar
comfyanonymous committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175

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

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "_for_testing"

    def encode(self, vae, pixels):
        x = (pixels.shape[1] // 64) * 64
        y = (pixels.shape[2] // 64) * 64
        if pixels.shape[1] != x or pixels.shape[2] != y:
            pixels = pixels[:,:x,:y,:]
        t = vae.encode_tiled(pixels[:,:,:,:3])

        return ({"samples":t}, )
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
class VAEEncodeForInpaint:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", ), "mask": ("MASK", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "latent/inpaint"

    def encode(self, vae, pixels, mask):
        x = (pixels.shape[1] // 64) * 64
        y = (pixels.shape[2] // 64) * 64
191
192
        mask = torch.nn.functional.interpolate(mask[None,None,], size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")[0][0]

193
        pixels = pixels.clone()
194
195
196
197
        if pixels.shape[1] != x or pixels.shape[2] != y:
            pixels = pixels[:,:x,:y,:]
            mask = mask[:x,:y]

198
        #grow mask by a few pixels to keep things seamless in latent space
199
        kernel_tensor = torch.ones((1, 1, 6, 6))
200
201
        mask_erosion = torch.clamp(torch.nn.functional.conv2d((mask.round())[None], kernel_tensor, padding=3), 0, 1)
        m = (1.0 - mask.round())
202
203
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
204
            pixels[:,:,:,i] *= m
205
206
207
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

208
        return ({"samples":t, "noise_mask": (mask_erosion[0][:x,:y].round())}, )
comfyanonymous's avatar
comfyanonymous committed
209
210
211

class CheckpointLoader:
    models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
212
    embedding_directory = os.path.join(models_dir, "embeddings")
comfyanonymous's avatar
comfyanonymous committed
213
214
215

    @classmethod
    def INPUT_TYPES(s):
216
217
        return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ),
                              "ckpt_name": (folder_paths.get_filename_list("checkpoints"), )}}
comfyanonymous's avatar
comfyanonymous committed
218
219
220
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

221
222
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
223
    def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
224
225
        config_path = folder_paths.get_full_path("configs", config_name)
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
226
        return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=self.embedding_directory)
comfyanonymous's avatar
comfyanonymous committed
227

228
229
230
class CheckpointLoaderSimple:
    @classmethod
    def INPUT_TYPES(s):
231
        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
232
233
234
235
                             }}
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

236
    CATEGORY = "loaders"
237

238
    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
239
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
240
        out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=CheckpointLoader.embedding_directory)
241
242
        return out

comfyanonymous's avatar
comfyanonymous committed
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
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,)

259
260
261
262
263
class LoraLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "clip": ("CLIP", ),
264
                              "lora_name": (folder_paths.get_filename_list("loras"), ),
265
266
267
268
269
270
271
272
273
                              "strength_model": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                              "strength_clip": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                              }}
    RETURN_TYPES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"

    CATEGORY = "loaders"

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
274
        lora_path = folder_paths.get_full_path("loras", lora_name)
275
276
277
        model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora_path, strength_model, strength_clip)
        return (model_lora, clip_lora)

comfyanonymous's avatar
comfyanonymous committed
278
279
280
class VAELoader:
    @classmethod
    def INPUT_TYPES(s):
281
        return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}}
comfyanonymous's avatar
comfyanonymous committed
282
283
284
    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"

285
286
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
287
288
    #TODO: scale factor?
    def load_vae(self, vae_name):
289
        vae_path = folder_paths.get_full_path("vae", vae_name)
comfyanonymous's avatar
comfyanonymous committed
290
291
292
        vae = comfy.sd.VAE(ckpt_path=vae_path)
        return (vae,)

comfyanonymous's avatar
comfyanonymous committed
293
294
295
class ControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
296
        return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
comfyanonymous's avatar
comfyanonymous committed
297
298
299
300
301
302
303

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

    CATEGORY = "loaders"

    def load_controlnet(self, control_net_name):
304
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
comfyanonymous's avatar
comfyanonymous committed
305
306
307
        controlnet = comfy.sd.load_controlnet(controlnet_path)
        return (controlnet,)

308
309
310
311
class DiffControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
312
                              "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
313
314
315
316
317
318
319

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

    CATEGORY = "loaders"

    def load_controlnet(self, model, control_net_name):
320
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
321
322
323
        controlnet = comfy.sd.load_controlnet(controlnet_path, model)
        return (controlnet,)

comfyanonymous's avatar
comfyanonymous committed
324
325
326
327

class ControlNetApply:
    @classmethod
    def INPUT_TYPES(s):
328
329
330
331
332
        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
333
334
335
336
337
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

338
    def apply_controlnet(self, conditioning, control_net, image, strength):
comfyanonymous's avatar
comfyanonymous committed
339
340
341
342
343
        c = []
        control_hint = image.movedim(-1,1)
        print(control_hint.shape)
        for t in conditioning:
            n = [t[0], t[1].copy()]
comfyanonymous's avatar
comfyanonymous committed
344
345
346
347
            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
348
349
350
            c.append(n)
        return (c, )

351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
class T2IAdapterLoader:
    models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
    t2i_adapter_dir = os.path.join(models_dir, "t2i_adapter")
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "t2i_adapter_name": (filter_files_extensions(recursive_search(s.t2i_adapter_dir), supported_pt_extensions), )}}

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

    CATEGORY = "loaders"

    def load_t2i_adapter(self, t2i_adapter_name):
        t2i_path = os.path.join(self.t2i_adapter_dir, t2i_adapter_name)
        t2i_adapter = comfy.sd.load_t2i_adapter(t2i_path)
        return (t2i_adapter,)
comfyanonymous's avatar
comfyanonymous committed
367

368
369
370
class CLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
371
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
372
373
374
375
376
377
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

378
    def load_clip(self, clip_name):
379
        clip_path = folder_paths.get_full_path("clip", clip_name)
380
381
382
        clip = comfy.sd.load_clip(ckpt_path=clip_path, embedding_directory=CheckpointLoader.embedding_directory)
        return (clip,)

383
384
385
class CLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
386
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
387
388
389
390
391
392
393
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

    def load_clip(self, clip_name):
394
        clip_path = folder_paths.get_full_path("clip_vision", clip_name)
395
396
397
398
399
400
401
402
403
        clip_vision = comfy_extras.clip_vision.load(clip_path)
        return (clip_vision,)

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

comfyanonymous's avatar
comfyanonymous committed
407
    CATEGORY = "conditioning/style_model"
408
409
410
411
412
413
414
415

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

class StyleModelLoader:
    @classmethod
    def INPUT_TYPES(s):
416
        return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}}
417
418
419
420
421
422
423

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

    CATEGORY = "loaders"

    def load_style_model(self, style_model_name):
424
        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
425
426
427
428
429
430
431
        style_model = comfy.sd.load_style_model(style_model_path)
        return (style_model,)


class StyleModelApply:
    @classmethod
    def INPUT_TYPES(s):
432
433
434
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "style_model": ("STYLE_MODEL", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
435
436
437
438
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_stylemodel"

comfyanonymous's avatar
comfyanonymous committed
439
    CATEGORY = "conditioning/style_model"
440

441
442
    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
        cond = style_model.get_cond(clip_vision_output)
443
        c = []
444
445
        for t in conditioning:
            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
446
447
448
            c.append(n)
        return (c, )

comfyanonymous's avatar
comfyanonymous committed
449
450
451
452
453
454
455
456
457
458
459
460
class EmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

461
462
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
463
464
    def generate(self, width, height, batch_size=1):
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
465
        return ({"samples":latent}, )
comfyanonymous's avatar
comfyanonymous committed
466

comfyanonymous's avatar
comfyanonymous committed
467

comfyanonymous's avatar
comfyanonymous committed
468

comfyanonymous's avatar
comfyanonymous committed
469
470
class LatentUpscale:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
471
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
472
473
474
475
476

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
                              "width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
477
478
                              "height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
479
480
481
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

482
483
    CATEGORY = "latent"

484
    def upscale(self, samples, upscale_method, width, height, crop):
485
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
486
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
487
488
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
489
490
491
492
493
494
495
496
497
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
498
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
499
500

    def rotate(self, samples, rotation):
501
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
502
503
504
505
506
507
508
509
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

510
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
511
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
512
513
514
515
516
517
518
519
520
521

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
522
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
523
524

    def flip(self, samples, flip_method):
525
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
526
        if flip_method.startswith("x"):
527
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
528
        elif flip_method.startswith("y"):
529
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
530
531

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
532
533
534
535
536
537
538
539

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples_to": ("LATENT",),
                              "samples_from": ("LATENT",),
                              "x": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
540
                              "feather": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
comfyanonymous's avatar
comfyanonymous committed
541
542
543
544
545
546
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

547
    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
comfyanonymous's avatar
comfyanonymous committed
548
549
        x =  x // 8
        y = y // 8
550
        feather = feather // 8
551
552
553
554
        samples_out = samples_to.copy()
        s = samples_to["samples"].clone()
        samples_to = samples_to["samples"]
        samples_from = samples_from["samples"]
555
556
557
        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:
558
559
            samples_from = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
            mask = torch.ones_like(samples_from)
560
561
562
563
564
565
566
567
568
569
570
571
            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
572
573
        samples_out["samples"] = s
        return (samples_out,)
comfyanonymous's avatar
comfyanonymous committed
574

comfyanonymous's avatar
comfyanonymous committed
575
576
577
578
579
580
581
582
583
584
585
586
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}),
                              "x": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                              "y": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
587
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
588
589

    def crop(self, samples, width, height, x, y):
590
591
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
        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
        def enforce_image_dim(d, to_d, max_d):
            if to_d > max_d:
                leftover = (to_d - max_d) % 8
                to_d = max_d
                d -= leftover
            return (d, to_d)

        #make sure size is always multiple of 64
        x, to_x = enforce_image_dim(x, to_x, samples.shape[3])
        y, to_y = enforce_image_dim(y, to_y, samples.shape[2])
615
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
616
617
        return (s,)

618
619
620
621
622
623
624
625
626
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

627
    CATEGORY = "latent/inpaint"
628
629
630
631
632
633
634

    def set_mask(self, samples, mask):
        s = samples.copy()
        s["noise_mask"] = mask
        return (s,)


635
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):
636
637
    latent_image = latent["samples"]
    noise_mask = None
638
    device = model_management.get_torch_device()
639

comfyanonymous's avatar
comfyanonymous committed
640
641
642
643
644
    if disable_noise:
        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
    else:
        noise = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=torch.manual_seed(seed), device="cpu")

645
646
647
    if "noise_mask" in latent:
        noise_mask = latent['noise_mask']
        noise_mask = torch.nn.functional.interpolate(noise_mask[None,None,], size=(noise.shape[2], noise.shape[3]), mode="bilinear")
648
        noise_mask = noise_mask.round()
649
650
651
652
        noise_mask = torch.cat([noise_mask] * noise.shape[1], dim=1)
        noise_mask = torch.cat([noise_mask] * noise.shape[0])
        noise_mask = noise_mask.to(device)

653
    real_model = None
654
655
656
    model_management.load_model_gpu(model)
    real_model = model.model

657
658
659
660
661
662
    noise = noise.to(device)
    latent_image = latent_image.to(device)

    positive_copy = []
    negative_copy = []

comfyanonymous's avatar
comfyanonymous committed
663
    control_nets = []
664
665
666
667
668
    for p in positive:
        t = p[0]
        if t.shape[0] < noise.shape[0]:
            t = torch.cat([t] * noise.shape[0])
        t = t.to(device)
comfyanonymous's avatar
comfyanonymous committed
669
670
        if 'control' in p[1]:
            control_nets += [p[1]['control']]
671
672
673
674
675
676
        positive_copy += [[t] + p[1:]]
    for n in negative:
        t = n[0]
        if t.shape[0] < noise.shape[0]:
            t = torch.cat([t] * noise.shape[0])
        t = t.to(device)
677
678
        if 'control' in n[1]:
            control_nets += [n[1]['control']]
679
680
        negative_copy += [[t] + n[1:]]

comfyanonymous's avatar
comfyanonymous committed
681
682
683
684
    control_net_models = []
    for x in control_nets:
        control_net_models += x.get_control_models()
    model_management.load_controlnet_gpu(control_net_models)
comfyanonymous's avatar
comfyanonymous committed
685

686
687
688
689
690
691
    if sampler_name in comfy.samplers.KSampler.SAMPLERS:
        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise)
    else:
        #other samplers
        pass

692
    samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask)
693
    samples = samples.cpu()
comfyanonymous's avatar
comfyanonymous committed
694
695
    for c in control_nets:
        c.cleanup()
comfyanonymous's avatar
comfyanonymous committed
696

697
698
699
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
700

comfyanonymous's avatar
comfyanonymous committed
701
702
703
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
704
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
                    {"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}),
                    }}

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

720
721
    CATEGORY = "sampling"

comfyanonymous's avatar
comfyanonymous committed
722
    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
723
        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
comfyanonymous's avatar
comfyanonymous committed
724

comfyanonymous's avatar
comfyanonymous committed
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
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"], ),
                    }}

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

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

comfyanonymous's avatar
comfyanonymous committed
749
750
751
752
753
754
755
    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):
        force_full_denoise = True
        if return_with_leftover_noise == "enable":
            force_full_denoise = False
        disable_noise = False
        if add_noise == "disable":
            disable_noise = True
756
        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
757
758
759
760

class SaveImage:
    def __init__(self):
        self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
pythongosssss's avatar
pythongosssss committed
761
        self.url_suffix = ""
comfyanonymous's avatar
comfyanonymous committed
762
763
764
765

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
766
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
767
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
768
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
769
770
771
772
773
774
775
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

776
777
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
778
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
779
780
781
782
783
784
785
786
787
788
789
790
        def map_filename(filename):
            prefix_len = len(filename_prefix)
            prefix = filename[:prefix_len + 1]
            try:
                digits = int(filename[prefix_len + 1:].split('_')[0])
            except:
                digits = 0
            return (digits, prefix)
        try:
            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(self.output_dir))))[0] + 1
        except ValueError:
            counter = 1
791
792
793
        except FileNotFoundError:
            os.mkdir(self.output_dir)
            counter = 1
pythongosssss's avatar
pythongosssss committed
794

pythongosssss's avatar
pythongosssss committed
795
796
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
797

pythongosssss's avatar
pythongosssss committed
798
        paths = list()
comfyanonymous's avatar
comfyanonymous committed
799
800
        for image in images:
            i = 255. * image.cpu().numpy()
801
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
comfyanonymous's avatar
comfyanonymous committed
802
803
804
805
806
807
            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]))
pythongosssss's avatar
pythongosssss committed
808
            file = f"{filename_prefix}_{counter:05}_.png"
pythongosssss's avatar
pythongosssss committed
809
810
            img.save(os.path.join(self.output_dir, file), pnginfo=metadata, optimize=True)
            paths.append(file + self.url_suffix)
811
            counter += 1
pythongosssss's avatar
pythongosssss committed
812
        return { "ui": { "images": paths } }
comfyanonymous's avatar
comfyanonymous committed
813

pythongosssss's avatar
pythongosssss committed
814
815
816
817
818
819
820
class PreviewImage(SaveImage):
    def __init__(self):
        self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
        self.url_suffix = "?type=temp"

    @classmethod
    def INPUT_TYPES(s):
821
        return {"required":
pythongosssss's avatar
pythongosssss committed
822
823
824
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
825

826
827
828
829
class LoadImage:
    input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
    @classmethod
    def INPUT_TYPES(s):
pythongosssss's avatar
pythongosssss committed
830
831
        if not os.path.exists(s.input_dir):
            os.makedirs(s.input_dir)
832
        return {"required":
833
                    {"image": (sorted(os.listdir(s.input_dir)), )},
834
                }
835
836

    CATEGORY = "image"
837

838
    RETURN_TYPES = ("IMAGE", "MASK")
839
840
841
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = os.path.join(self.input_dir, image)
842
843
        i = Image.open(image_path)
        image = i.convert("RGB")
844
        image = np.array(image).astype(np.float32) / 255.0
845
        image = torch.from_numpy(image)[None,]
846
847
848
849
850
851
        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)
852

853
854
855
856
857
858
859
    @classmethod
    def IS_CHANGED(s, image):
        image_path = os.path.join(s.input_dir, image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
860

861
862
863
864
865
class LoadImageMask:
    input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
866
                    {"image": (sorted(os.listdir(s.input_dir)), ),
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
                    "channel": (["alpha", "red", "green", "blue"], ),}
                }

    CATEGORY = "image"

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
        image_path = os.path.join(self.input_dir, image)
        i = Image.open(image_path)
        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):
        image_path = os.path.join(s.input_dir, image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
895

comfyanonymous's avatar
comfyanonymous committed
896
897
898
899
900
901
902
903
904
905
906
907
908
class ImageScale:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
    crop_methods = ["disabled", "center"]

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
                              "width": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1}),
                              "height": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1}),
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

909
    CATEGORY = "image/upscaling"
910

comfyanonymous's avatar
comfyanonymous committed
911
912
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
913
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
914
915
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
916

917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
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,)


comfyanonymous's avatar
comfyanonymous committed
933
934
935
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
    "CheckpointLoader": CheckpointLoader,
936
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
937
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
938
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
939
940
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
941
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
942
943
944
945
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
    "LatentUpscale": LatentUpscale,
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
946
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
947
    "LoadImage": LoadImage,
948
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
949
    "ImageScale": ImageScale,
950
    "ImageInvert": ImageInvert,
comfyanonymous's avatar
comfyanonymous committed
951
952
    "ConditioningCombine": ConditioningCombine,
    "ConditioningSetArea": ConditioningSetArea,
comfyanonymous's avatar
comfyanonymous committed
953
    "KSamplerAdvanced": KSamplerAdvanced,
954
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
955
    "LatentComposite": LatentComposite,
comfyanonymous's avatar
comfyanonymous committed
956
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
957
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
958
    "LatentCrop": LatentCrop,
959
    "LoraLoader": LoraLoader,
960
    "CLIPLoader": CLIPLoader,
961
    "CLIPVisionEncode": CLIPVisionEncode,
962
    "StyleModelApply": StyleModelApply,
comfyanonymous's avatar
comfyanonymous committed
963
964
    "ControlNetApply": ControlNetApply,
    "ControlNetLoader": ControlNetLoader,
965
    "DiffControlNetLoader": DiffControlNetLoader,
966
    "T2IAdapterLoader": T2IAdapterLoader,
comfyanonymous's avatar
comfyanonymous committed
967
968
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
969
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
970
    "VAEEncodeTiled": VAEEncodeTiled,
comfyanonymous's avatar
comfyanonymous committed
971
972
}

973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
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)
        else:
            print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
    except Exception as e:
        print(traceback.format_exc())
        print(f"Cannot import {module_path} module for custom nodes:", e)

Hacker 17082006's avatar
Hacker 17082006 committed
994
def load_custom_nodes():
995
    CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
996
    possible_modules = os.listdir(CUSTOM_NODE_PATH)
997
    if "__pycache__" in possible_modules:
Hacker 17082006's avatar
.  
Hacker 17082006 committed
998
        possible_modules.remove("__pycache__")
999

Hacker 17082006's avatar
Hacker 17082006 committed
1000
    for possible_module in possible_modules:
1001
1002
        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
1003
        load_custom_node(module_path)
1004

1005
load_custom_nodes()
1006
1007

load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))