nodes.py 40.2 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
comfyanonymous's avatar
comfyanonymous committed
8
9
10
11
12

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

sALTaccount's avatar
sALTaccount committed
13
14
from comfy.diffusers_convert import load_diffusers

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


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

22
import comfy.clip_vision
23

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

27
import folder_paths
28
29
30
31

def before_node_execution():
    model_management.throw_exception_if_processing_interrupted()

32
33
def interrupt_processing(value=True):
    model_management.interrupt_current_processing(value)
34

35
36
MAX_RESOLUTION=8192

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

44
45
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
46
    def encode(self, clip, text):
comfyanonymous's avatar
comfyanonymous committed
47
48
49
50
51
52
53
54
55
        return ([[clip.encode(text), {}]], )

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

56
57
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
58
59
60
61
62
63
64
    def combine(self, conditioning_1, conditioning_2):
        return (conditioning_1 + conditioning_2, )

class ConditioningSetArea:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
65
66
67
68
                              "width": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
comfyanonymous's avatar
comfyanonymous committed
69
70
71
72
73
                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "append"

74
75
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
76
    def append(self, conditioning, width, height, x, y, strength, min_sigma=0.0, max_sigma=99.0):
comfyanonymous's avatar
comfyanonymous committed
77
78
79
80
81
82
83
84
        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
85
        return (c, )
comfyanonymous's avatar
comfyanonymous committed
86
87
88
89
90
91
92
93
94
95
96

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"

97
98
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
99
    def decode(self, vae, samples):
100
        return (vae.decode(samples["samples"]), )
comfyanonymous's avatar
comfyanonymous committed
101

102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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
117
118
119
120
121
122
123
124
125
126
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"

127
128
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
129
    def encode(self, vae, pixels):
130
131
132
133
        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,:]
134
135
136
        t = vae.encode(pixels[:,:,:,:3])

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

comfyanonymous's avatar
comfyanonymous committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

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}, )
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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
174
175
        mask = torch.nn.functional.interpolate(mask[None,None,], size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")[0][0]

176
        pixels = pixels.clone()
177
178
179
180
        if pixels.shape[1] != x or pixels.shape[2] != y:
            pixels = pixels[:,:x,:y,:]
            mask = mask[:x,:y]

181
        #grow mask by a few pixels to keep things seamless in latent space
182
        kernel_tensor = torch.ones((1, 1, 6, 6))
183
184
        mask_erosion = torch.clamp(torch.nn.functional.conv2d((mask.round())[None], kernel_tensor, padding=3), 0, 1)
        m = (1.0 - mask.round())
185
186
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
187
            pixels[:,:,:,i] *= m
188
189
190
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

191
        return ({"samples":t, "noise_mask": (mask_erosion[0][:x,:y].round())}, )
comfyanonymous's avatar
comfyanonymous committed
192
193
194
195

class CheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
196
197
        return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ),
                              "ckpt_name": (folder_paths.get_filename_list("checkpoints"), )}}
comfyanonymous's avatar
comfyanonymous committed
198
199
200
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

201
    CATEGORY = "advanced/loaders"
202

comfyanonymous's avatar
comfyanonymous committed
203
    def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
204
205
        config_path = folder_paths.get_full_path("configs", config_name)
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
206
        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
207

208
209
210
class CheckpointLoaderSimple:
    @classmethod
    def INPUT_TYPES(s):
211
        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
212
213
214
215
                             }}
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

216
    CATEGORY = "loaders"
217

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

sALTaccount's avatar
sALTaccount committed
223
224
225
class DiffusersLoader:
    @classmethod
    def INPUT_TYPES(cls):
226
227
228
229
230
        paths = []
        search_path = os.path.join(folder_paths.models_dir, 'diffusers')
        if os.path.exists(search_path):
            paths = next(os.walk(search_path))[1]
        return {"required": {"model_path": (paths,), }}
sALTaccount's avatar
sALTaccount committed
231
232
233
234
235
236
237
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

    CATEGORY = "loaders"

    def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
        model_path = os.path.join(folder_paths.models_dir, 'diffusers', model_path)
238
        return load_diffusers(model_path, fp16=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
239
240


241
242
243
244
245
246
247
248
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"

249
    CATEGORY = "loaders"
250
251
252
253
254
255

    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
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
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,)

272
273
274
275
276
class LoraLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "clip": ("CLIP", ),
277
                              "lora_name": (folder_paths.get_filename_list("loras"), ),
278
279
                              "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}),
280
281
282
283
284
285
286
                              }}
    RETURN_TYPES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"

    CATEGORY = "loaders"

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
287
        lora_path = folder_paths.get_full_path("loras", lora_name)
288
289
290
        model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora_path, strength_model, strength_clip)
        return (model_lora, clip_lora)

291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
class TomePatchModel:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
                              }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "patch"

    CATEGORY = "_for_testing"

    def patch(self, model, ratio):
        m = model.clone()
        m.set_model_tomesd(ratio)
        return (m, )

comfyanonymous's avatar
comfyanonymous committed
307
308
309
class VAELoader:
    @classmethod
    def INPUT_TYPES(s):
310
        return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}}
comfyanonymous's avatar
comfyanonymous committed
311
312
313
    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"

314
315
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
316
317
    #TODO: scale factor?
    def load_vae(self, vae_name):
318
        vae_path = folder_paths.get_full_path("vae", vae_name)
comfyanonymous's avatar
comfyanonymous committed
319
320
321
        vae = comfy.sd.VAE(ckpt_path=vae_path)
        return (vae,)

comfyanonymous's avatar
comfyanonymous committed
322
323
324
class ControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
325
        return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
comfyanonymous's avatar
comfyanonymous committed
326
327
328
329
330
331
332

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

    CATEGORY = "loaders"

    def load_controlnet(self, control_net_name):
333
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
comfyanonymous's avatar
comfyanonymous committed
334
335
336
        controlnet = comfy.sd.load_controlnet(controlnet_path)
        return (controlnet,)

337
338
339
340
class DiffControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
341
                              "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
342
343
344
345
346
347
348

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

    CATEGORY = "loaders"

    def load_controlnet(self, model, control_net_name):
349
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
350
351
352
        controlnet = comfy.sd.load_controlnet(controlnet_path, model)
        return (controlnet,)

comfyanonymous's avatar
comfyanonymous committed
353
354
355
356

class ControlNetApply:
    @classmethod
    def INPUT_TYPES(s):
357
358
359
360
361
        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
362
363
364
365
366
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

367
    def apply_controlnet(self, conditioning, control_net, image, strength):
comfyanonymous's avatar
comfyanonymous committed
368
369
370
371
372
        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
373
374
375
376
            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
377
378
379
            c.append(n)
        return (c, )

380
381
382
class CLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
383
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
384
385
386
387
388
389
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

390
    def load_clip(self, clip_name):
391
        clip_path = folder_paths.get_full_path("clip", clip_name)
comfyanonymous's avatar
comfyanonymous committed
392
        clip = comfy.sd.load_clip(ckpt_path=clip_path, embedding_directory=folder_paths.get_folder_paths("embeddings"))
393
394
        return (clip,)

395
396
397
class CLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
398
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
399
400
401
402
403
404
405
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

    def load_clip(self, clip_name):
406
        clip_path = folder_paths.get_full_path("clip_vision", clip_name)
407
        clip_vision = comfy.clip_vision.load(clip_path)
408
409
410
411
412
413
414
415
        return (clip_vision,)

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

419
    CATEGORY = "conditioning"
420
421
422
423
424
425
426
427

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

class StyleModelLoader:
    @classmethod
    def INPUT_TYPES(s):
428
        return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}}
429
430
431
432
433
434
435

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

    CATEGORY = "loaders"

    def load_style_model(self, style_model_name):
436
        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
437
438
439
440
441
442
443
        style_model = comfy.sd.load_style_model(style_model_path)
        return (style_model,)


class StyleModelApply:
    @classmethod
    def INPUT_TYPES(s):
444
445
446
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "style_model": ("STYLE_MODEL", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
447
448
449
450
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_stylemodel"

comfyanonymous's avatar
comfyanonymous committed
451
    CATEGORY = "conditioning/style_model"
452

453
454
    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
        cond = style_model.get_cond(clip_vision_output)
455
        c = []
456
457
        for t in conditioning:
            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
458
459
460
            c.append(n)
        return (c, )

461
462
463
464
465
466
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}),
467
                             "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
468
469
470
471
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_adm"

472
    CATEGORY = "conditioning"
473

474
    def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation):
475
476
477
        c = []
        for t in conditioning:
            o = t[1].copy()
478
            x = (clip_vision_output, strength, noise_augmentation)
479
480
481
482
483
484
485
486
487
            if "adm" in o:
                o["adm"] = o["adm"][:] + [x]
            else:
                o["adm"] = [x]
            n = [t[0], o]
            c.append(n)
        return (c, )


comfyanonymous's avatar
comfyanonymous committed
488
489
490
491
492
493
class EmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
494
495
        return {"required": { "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
comfyanonymous's avatar
comfyanonymous committed
496
497
498
499
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

500
501
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
502
503
    def generate(self, width, height, batch_size=1):
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
504
        return ({"samples":latent}, )
comfyanonymous's avatar
comfyanonymous committed
505

comfyanonymous's avatar
comfyanonymous committed
506

comfyanonymous's avatar
comfyanonymous committed
507

comfyanonymous's avatar
comfyanonymous committed
508
509
class LatentUpscale:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
510
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
511
512
513
514

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
515
516
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
517
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
518
519
520
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

521
522
    CATEGORY = "latent"

523
    def upscale(self, samples, upscale_method, width, height, crop):
524
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
525
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
526
527
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
528
529
530
531
532
533
534
535
536
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
537
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
538
539

    def rotate(self, samples, rotation):
540
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
541
542
543
544
545
546
547
548
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

549
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
550
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
551
552
553
554
555
556
557
558
559
560

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
561
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
562
563

    def flip(self, samples, flip_method):
564
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
565
        if flip_method.startswith("x"):
566
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
567
        elif flip_method.startswith("y"):
568
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
569
570

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
571
572
573
574
575
576

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples_to": ("LATENT",),
                              "samples_from": ("LATENT",),
577
578
579
                              "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
580
581
582
583
584
585
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

586
    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
comfyanonymous's avatar
comfyanonymous committed
587
588
        x =  x // 8
        y = y // 8
589
        feather = feather // 8
590
591
592
593
        samples_out = samples_to.copy()
        s = samples_to["samples"].clone()
        samples_to = samples_to["samples"]
        samples_from = samples_from["samples"]
594
595
596
        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:
597
598
            samples_from = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
            mask = torch.ones_like(samples_from)
599
600
601
602
603
604
605
606
607
608
609
610
            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
611
612
        samples_out["samples"] = s
        return (samples_out,)
comfyanonymous's avatar
comfyanonymous committed
613

comfyanonymous's avatar
comfyanonymous committed
614
615
616
617
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
618
619
620
621
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "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
622
623
624
625
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
626
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
627
628

    def crop(self, samples, width, height, x, y):
629
630
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
        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])
654
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
655
656
        return (s,)

657
658
659
660
661
662
663
664
665
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

666
    CATEGORY = "latent/inpaint"
667
668
669
670
671
672
673

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


674
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):
675
676
    latent_image = latent["samples"]
    noise_mask = None
677
    device = model_management.get_torch_device()
678

comfyanonymous's avatar
comfyanonymous committed
679
680
681
682
683
    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")

684
685
686
    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")
687
        noise_mask = noise_mask.round()
688
689
690
691
        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)

692
    real_model = None
693
694
695
    model_management.load_model_gpu(model)
    real_model = model.model

696
697
698
699
700
701
    noise = noise.to(device)
    latent_image = latent_image.to(device)

    positive_copy = []
    negative_copy = []

comfyanonymous's avatar
comfyanonymous committed
702
    control_nets = []
703
704
705
706
707
    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
708
709
        if 'control' in p[1]:
            control_nets += [p[1]['control']]
710
711
712
713
714
715
        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)
716
717
        if 'control' in n[1]:
            control_nets += [n[1]['control']]
718
719
        negative_copy += [[t] + n[1:]]

comfyanonymous's avatar
comfyanonymous committed
720
721
722
723
    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
724

725
    if sampler_name in comfy.samplers.KSampler.SAMPLERS:
726
        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
727
728
729
730
    else:
        #other samplers
        pass

731
    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)
732
    samples = samples.cpu()
comfyanonymous's avatar
comfyanonymous committed
733
734
    for c in control_nets:
        c.cleanup()
comfyanonymous's avatar
comfyanonymous committed
735

736
737
738
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
739

comfyanonymous's avatar
comfyanonymous committed
740
741
742
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
743
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
                    {"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"

759
760
    CATEGORY = "sampling"

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

comfyanonymous's avatar
comfyanonymous committed
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
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
787

comfyanonymous's avatar
comfyanonymous committed
788
789
790
791
792
793
794
    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
795
        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
796
797
798

class SaveImage:
    def __init__(self):
799
        self.output_dir = folder_paths.get_output_directory()
m957ymj75urz's avatar
m957ymj75urz committed
800
        self.type = "output"
comfyanonymous's avatar
comfyanonymous committed
801
802
803
804

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
805
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
806
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
807
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
808
809
810
811
812
813
814
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

815
816
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
817
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
818
        def map_filename(filename):
819
            prefix_len = len(os.path.basename(filename_prefix))
820
821
822
823
824
825
            prefix = filename[:prefix_len + 1]
            try:
                digits = int(filename[prefix_len + 1:].split('_')[0])
            except:
                digits = 0
            return (digits, prefix)
comfyanonymous's avatar
Style.  
comfyanonymous committed
826

827
828
829
830
        def compute_vars(input):
            input = input.replace("%width%", str(images[0].shape[1]))
            input = input.replace("%height%", str(images[0].shape[0]))
            return input
comfyanonymous's avatar
Style.  
comfyanonymous committed
831

832
        filename_prefix = compute_vars(filename_prefix)
comfyanonymous's avatar
comfyanonymous committed
833

m957ymj75urz's avatar
m957ymj75urz committed
834
835
836
        subfolder = os.path.dirname(os.path.normpath(filename_prefix))
        filename = os.path.basename(os.path.normpath(filename_prefix))

comfyanonymous's avatar
comfyanonymous committed
837
        full_output_folder = os.path.join(self.output_dir, subfolder)
838

839
        if os.path.commonpath((self.output_dir, os.path.abspath(full_output_folder))) != self.output_dir:
840
            print("Saving image outside the output folder is not allowed.")
comfyanonymous's avatar
comfyanonymous committed
841
842
            return {}

843
        try:
844
            counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
845
846
        except ValueError:
            counter = 1
847
        except FileNotFoundError:
848
            os.makedirs(full_output_folder, exist_ok=True)
849
            counter = 1
pythongosssss's avatar
pythongosssss committed
850

m957ymj75urz's avatar
m957ymj75urz committed
851
        results = list()
comfyanonymous's avatar
comfyanonymous committed
852
853
        for image in images:
            i = 255. * image.cpu().numpy()
854
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
comfyanonymous's avatar
comfyanonymous committed
855
856
857
858
859
860
            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]))
861

862
            file = f"{filename}_{counter:05}_.png"
863
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
m957ymj75urz's avatar
m957ymj75urz committed
864
865
866
867
868
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
            });
869
            counter += 1
comfyanonymous's avatar
comfyanonymous committed
870

m957ymj75urz's avatar
m957ymj75urz committed
871
        return { "ui": { "images": results } }
comfyanonymous's avatar
comfyanonymous committed
872

pythongosssss's avatar
pythongosssss committed
873
874
class PreviewImage(SaveImage):
    def __init__(self):
875
        self.output_dir = folder_paths.get_temp_directory()
m957ymj75urz's avatar
m957ymj75urz committed
876
        self.type = "temp"
pythongosssss's avatar
pythongosssss committed
877
878
879

    @classmethod
    def INPUT_TYPES(s):
880
        return {"required":
pythongosssss's avatar
pythongosssss committed
881
882
883
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
884

885
886
887
class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
888
        input_dir = folder_paths.get_input_directory()
889
        return {"required":
890
                    {"image": (sorted(os.listdir(input_dir)), )},
891
                }
892
893

    CATEGORY = "image"
894

895
    RETURN_TYPES = ("IMAGE", "MASK")
896
897
    FUNCTION = "load_image"
    def load_image(self, image):
898
899
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
900
901
        i = Image.open(image_path)
        image = i.convert("RGB")
902
        image = np.array(image).astype(np.float32) / 255.0
903
        image = torch.from_numpy(image)[None,]
904
905
906
907
908
909
        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)
910

911
912
    @classmethod
    def IS_CHANGED(s, image):
913
914
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
915
916
917
918
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
919

920
921
922
class LoadImageMask:
    @classmethod
    def INPUT_TYPES(s):
923
        input_dir = folder_paths.get_input_directory()
924
        return {"required":
925
                    {"image": (sorted(os.listdir(input_dir)), ),
926
927
928
929
930
931
932
933
                    "channel": (["alpha", "red", "green", "blue"], ),}
                }

    CATEGORY = "image"

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
934
935
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
936
937
938
939
940
941
942
943
944
945
946
947
948
949
        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):
950
951
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
952
953
954
955
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
956

comfyanonymous's avatar
comfyanonymous committed
957
958
959
960
961
962
963
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,),
964
965
                              "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
966
967
968
969
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

970
    CATEGORY = "image/upscaling"
971

comfyanonymous's avatar
comfyanonymous committed
972
973
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
974
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
975
976
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
977

978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
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
994
995
996
997
998
999
1000
class ImagePadForOutpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
1001
1002
1003
1004
1005
                "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
                "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
                "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
                "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}),
                "feathering": ("INT", {"default": 40, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
Guo Y.K's avatar
Guo Y.K committed
1006
1007
1008
1009
1010
1011
1012
1013
            }
        }

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

    CATEGORY = "image"

1014
    def expand_image(self, image, left, top, right, bottom, feathering):
Guo Y.K's avatar
Guo Y.K committed
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
        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,
        )
1027

1028
1029
1030
1031
1032
        t = torch.zeros(
            (d2, d3),
            dtype=torch.float32
        )

1033
        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052

            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
1053

Guo Y.K's avatar
Guo Y.K committed
1054
1055
1056
        return (new_image, mask)


comfyanonymous's avatar
comfyanonymous committed
1057
1058
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
1059
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
1060
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
1061
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
1062
1063
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
1064
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
1065
1066
1067
1068
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
    "LatentUpscale": LatentUpscale,
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
1069
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
1070
    "LoadImage": LoadImage,
1071
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
1072
    "ImageScale": ImageScale,
1073
    "ImageInvert": ImageInvert,
Guo Y.K's avatar
Guo Y.K committed
1074
    "ImagePadForOutpaint": ImagePadForOutpaint,
comfyanonymous's avatar
comfyanonymous committed
1075
1076
    "ConditioningCombine": ConditioningCombine,
    "ConditioningSetArea": ConditioningSetArea,
comfyanonymous's avatar
comfyanonymous committed
1077
    "KSamplerAdvanced": KSamplerAdvanced,
1078
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
1079
    "LatentComposite": LatentComposite,
comfyanonymous's avatar
comfyanonymous committed
1080
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
1081
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
1082
    "LatentCrop": LatentCrop,
1083
    "LoraLoader": LoraLoader,
1084
    "CLIPLoader": CLIPLoader,
1085
    "CLIPVisionEncode": CLIPVisionEncode,
1086
    "StyleModelApply": StyleModelApply,
1087
    "unCLIPConditioning": unCLIPConditioning,
comfyanonymous's avatar
comfyanonymous committed
1088
1089
    "ControlNetApply": ControlNetApply,
    "ControlNetLoader": ControlNetLoader,
1090
    "DiffControlNetLoader": DiffControlNetLoader,
comfyanonymous's avatar
comfyanonymous committed
1091
1092
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
1093
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
1094
    "VAEEncodeTiled": VAEEncodeTiled,
1095
    "TomePatchModel": TomePatchModel,
1096
    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
1097
    "CheckpointLoader": CheckpointLoader,
sALTaccount's avatar
sALTaccount committed
1098
    "DiffusersLoader": DiffusersLoader,
comfyanonymous's avatar
comfyanonymous committed
1099
1100
}

1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
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
1122
def load_custom_nodes():
1123
    CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
1124
    possible_modules = os.listdir(CUSTOM_NODE_PATH)
1125
    if "__pycache__" in possible_modules:
Hacker 17082006's avatar
.  
Hacker 17082006 committed
1126
        possible_modules.remove("__pycache__")
1127

Hacker 17082006's avatar
Hacker 17082006 committed
1128
    for possible_module in possible_modules:
1129
1130
        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
1131
        load_custom_node(module_path)
1132

1133
1134
def init_custom_nodes():
    load_custom_nodes()
1135
1136
    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"))