nodes.py 42.6 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

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


comfyanonymous's avatar
comfyanonymous committed
17
import comfy.diffusers_convert
comfyanonymous's avatar
comfyanonymous committed
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
        paths = []
sALTaccount's avatar
sALTaccount committed
227
        for search_path in folder_paths.get_folder_paths("diffusers"):
228
            if os.path.exists(search_path):
sALTaccount's avatar
sALTaccount committed
229
                paths += next(os.walk(search_path))[1]
230
        return {"required": {"model_path": (paths,), }}
sALTaccount's avatar
sALTaccount committed
231
232
233
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

234
    CATEGORY = "advanced/loaders"
sALTaccount's avatar
sALTaccount committed
235
236

    def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
sALTaccount's avatar
sALTaccount committed
237
238
239
240
241
242
        for search_path in folder_paths.get_folder_paths("diffusers"):
            if os.path.exists(search_path):
                paths = next(os.walk(search_path))[1]
                if model_path in paths:
                    model_path = os.path.join(search_path, model_path)
                    break
243

comfyanonymous's avatar
comfyanonymous committed
244
        return comfy.diffusers_convert.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
245
246


247
248
249
250
251
252
253
254
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"

255
    CATEGORY = "loaders"
256
257
258
259
260
261

    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
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
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,)

278
279
280
281
282
class LoraLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "clip": ("CLIP", ),
283
                              "lora_name": (folder_paths.get_filename_list("loras"), ),
284
285
                              "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}),
286
287
288
289
290
291
292
                              }}
    RETURN_TYPES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"

    CATEGORY = "loaders"

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
293
        lora_path = folder_paths.get_full_path("loras", lora_name)
294
295
296
        model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora_path, strength_model, strength_clip)
        return (model_lora, clip_lora)

297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
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
313
314
315
class VAELoader:
    @classmethod
    def INPUT_TYPES(s):
316
        return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}}
comfyanonymous's avatar
comfyanonymous committed
317
318
319
    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"

320
321
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
322
323
    #TODO: scale factor?
    def load_vae(self, vae_name):
324
        vae_path = folder_paths.get_full_path("vae", vae_name)
comfyanonymous's avatar
comfyanonymous committed
325
326
327
        vae = comfy.sd.VAE(ckpt_path=vae_path)
        return (vae,)

comfyanonymous's avatar
comfyanonymous committed
328
329
330
class ControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
331
        return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
comfyanonymous's avatar
comfyanonymous committed
332
333
334
335
336
337
338

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

    CATEGORY = "loaders"

    def load_controlnet(self, control_net_name):
339
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
comfyanonymous's avatar
comfyanonymous committed
340
341
342
        controlnet = comfy.sd.load_controlnet(controlnet_path)
        return (controlnet,)

343
344
345
346
class DiffControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
347
                              "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
348
349
350
351
352
353
354

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

    CATEGORY = "loaders"

    def load_controlnet(self, model, control_net_name):
355
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
356
357
358
        controlnet = comfy.sd.load_controlnet(controlnet_path, model)
        return (controlnet,)

comfyanonymous's avatar
comfyanonymous committed
359
360
361
362

class ControlNetApply:
    @classmethod
    def INPUT_TYPES(s):
363
364
365
366
367
        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
368
369
370
371
372
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

373
    def apply_controlnet(self, conditioning, control_net, image, strength):
comfyanonymous's avatar
comfyanonymous committed
374
375
376
377
378
        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
379
380
381
382
            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
383
384
385
            c.append(n)
        return (c, )

386
387
388
class CLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
389
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
390
391
392
393
394
395
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

396
    def load_clip(self, clip_name):
397
        clip_path = folder_paths.get_full_path("clip", clip_name)
comfyanonymous's avatar
comfyanonymous committed
398
        clip = comfy.sd.load_clip(ckpt_path=clip_path, embedding_directory=folder_paths.get_folder_paths("embeddings"))
399
400
        return (clip,)

401
402
403
class CLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
404
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
405
406
407
408
409
410
411
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

    def load_clip(self, clip_name):
412
        clip_path = folder_paths.get_full_path("clip_vision", clip_name)
413
        clip_vision = comfy.clip_vision.load(clip_path)
414
415
416
417
418
419
420
421
        return (clip_vision,)

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

425
    CATEGORY = "conditioning"
426
427
428
429
430
431
432
433

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

class StyleModelLoader:
    @classmethod
    def INPUT_TYPES(s):
434
        return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}}
435
436
437
438
439
440
441

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

    CATEGORY = "loaders"

    def load_style_model(self, style_model_name):
442
        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
443
444
445
446
447
448
449
        style_model = comfy.sd.load_style_model(style_model_path)
        return (style_model,)


class StyleModelApply:
    @classmethod
    def INPUT_TYPES(s):
450
451
452
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "style_model": ("STYLE_MODEL", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
453
454
455
456
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_stylemodel"

comfyanonymous's avatar
comfyanonymous committed
457
    CATEGORY = "conditioning/style_model"
458

459
460
    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
        cond = style_model.get_cond(clip_vision_output)
461
        c = []
462
463
        for t in conditioning:
            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
464
465
466
            c.append(n)
        return (c, )

467
468
469
470
471
472
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}),
473
                             "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
474
475
476
477
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_adm"

478
    CATEGORY = "conditioning"
479

480
    def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation):
481
482
483
        c = []
        for t in conditioning:
            o = t[1].copy()
484
            x = (clip_vision_output, strength, noise_augmentation)
485
486
487
488
489
490
491
492
493
            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
494
495
496
497
498
499
class EmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
500
501
        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
502
503
504
505
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

506
507
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
508
509
    def generate(self, width, height, batch_size=1):
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
510
        return ({"samples":latent}, )
comfyanonymous's avatar
comfyanonymous committed
511

comfyanonymous's avatar
comfyanonymous committed
512

comfyanonymous's avatar
comfyanonymous committed
513

comfyanonymous's avatar
comfyanonymous committed
514
515
class LatentUpscale:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
516
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
517
518
519
520

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
521
522
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
523
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
524
525
526
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

527
528
    CATEGORY = "latent"

529
    def upscale(self, samples, upscale_method, width, height, crop):
530
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
531
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
532
533
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
534
535
536
537
538
539
540
541
542
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
543
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
544
545

    def rotate(self, samples, rotation):
546
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
547
548
549
550
551
552
553
554
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

555
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
556
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
557
558
559
560
561
562
563
564
565
566

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
567
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
568
569

    def flip(self, samples, flip_method):
570
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
571
        if flip_method.startswith("x"):
572
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
573
        elif flip_method.startswith("y"):
574
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
575
576

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
577
578
579
580

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
581
582
583
584
585
586
587
588
589
        return {
            "required": {
                "samples_to": ("LATENT",),
                "samples_from": ("LATENT",),
                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "feather": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
            }
        }
comfyanonymous's avatar
comfyanonymous committed
590
591
592
593
594
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

595
596
597
598
599
600
601
    def composite(self, samples_to, samples_from, x, y, feather):
        output = samples_to.copy()
        destination = samples_to["samples"].clone()
        source = samples_from["samples"]

        left, top = (x // 8, y // 8)
        right, bottom = (left + source.shape[3], top + source.shape[2],)
602
        feather = feather // 8
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638



        # calculate the bounds of the source that will be overlapping the destination
        # this prevents the source trying to overwrite latent pixels that are out of bounds
        # of the destination
        visible_width, visible_height = (destination.shape[3] - left, destination.shape[2] - top,)

        mask = torch.ones_like(source)

        for f in range(feather):
            feather_rate = (f + 1.0) / feather

            if left > 0:
                mask[:, :, :, f] *= feather_rate

            if right < destination.shape[3] - 1:
                mask[:, :, :, -f] *= feather_rate

            if top > 0:
                mask[:, :, f, :] *= feather_rate

            if bottom < destination.shape[2] - 1:
                mask[:, :, -f, :] *= feather_rate

        mask = mask[:, :, :visible_height, :visible_width]
        inverse_mask = torch.ones_like(mask) - mask

        source_portion = mask * source[:, :, :visible_height, :visible_width]
        destination_portion = inverse_mask  * destination[:, :, top:bottom, left:right]

        destination[:, :, top:bottom, left:right] = source_portion + destination_portion

        output["samples"] = destination

        return (output,)
comfyanonymous's avatar
comfyanonymous committed
639

comfyanonymous's avatar
comfyanonymous committed
640
641
642
643
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
644
645
646
647
                              "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
648
649
650
651
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
652
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
653
654

    def crop(self, samples, width, height, x, y):
655
656
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
        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])
680
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
681
682
        return (s,)

683
684
685
686
687
688
689
690
691
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

692
    CATEGORY = "latent/inpaint"
693
694
695
696
697
698
699

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


700
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):
701
702
    latent_image = latent["samples"]
    noise_mask = None
703
    device = model_management.get_torch_device()
704

comfyanonymous's avatar
comfyanonymous committed
705
706
707
708
709
    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")

710
711
712
    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")
713
        noise_mask = noise_mask.round()
714
715
716
717
        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)

718
    real_model = None
719
720
721
    model_management.load_model_gpu(model)
    real_model = model.model

722
723
724
725
726
727
    noise = noise.to(device)
    latent_image = latent_image.to(device)

    positive_copy = []
    negative_copy = []

comfyanonymous's avatar
comfyanonymous committed
728
    control_nets = []
729
730
731
732
733
    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
734
735
        if 'control' in p[1]:
            control_nets += [p[1]['control']]
736
737
738
739
740
741
        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)
742
743
        if 'control' in n[1]:
            control_nets += [n[1]['control']]
744
745
        negative_copy += [[t] + n[1:]]

comfyanonymous's avatar
comfyanonymous committed
746
747
748
749
    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
750

751
    if sampler_name in comfy.samplers.KSampler.SAMPLERS:
752
        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
753
754
755
756
    else:
        #other samplers
        pass

757
    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)
758
    samples = samples.cpu()
comfyanonymous's avatar
comfyanonymous committed
759
760
    for c in control_nets:
        c.cleanup()
comfyanonymous's avatar
comfyanonymous committed
761

762
763
764
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
765

comfyanonymous's avatar
comfyanonymous committed
766
767
768
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
769
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
                    {"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"

785
786
    CATEGORY = "sampling"

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

comfyanonymous's avatar
comfyanonymous committed
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
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
813

comfyanonymous's avatar
comfyanonymous committed
814
815
816
817
818
819
820
    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
821
        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
822
823
824

class SaveImage:
    def __init__(self):
825
        self.output_dir = folder_paths.get_output_directory()
m957ymj75urz's avatar
m957ymj75urz committed
826
        self.type = "output"
comfyanonymous's avatar
comfyanonymous committed
827
828
829
830

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
831
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
832
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
833
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
834
835
836
837
838
839
840
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

841
842
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
843
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
844
        def map_filename(filename):
845
            prefix_len = len(os.path.basename(filename_prefix))
846
847
848
849
850
851
            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
852

853
854
855
856
        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
857

858
        filename_prefix = compute_vars(filename_prefix)
comfyanonymous's avatar
comfyanonymous committed
859

m957ymj75urz's avatar
m957ymj75urz committed
860
861
862
        subfolder = os.path.dirname(os.path.normpath(filename_prefix))
        filename = os.path.basename(os.path.normpath(filename_prefix))

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

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

869
        try:
870
            counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
871
872
        except ValueError:
            counter = 1
873
        except FileNotFoundError:
874
            os.makedirs(full_output_folder, exist_ok=True)
875
            counter = 1
pythongosssss's avatar
pythongosssss committed
876

m957ymj75urz's avatar
m957ymj75urz committed
877
        results = list()
comfyanonymous's avatar
comfyanonymous committed
878
879
        for image in images:
            i = 255. * image.cpu().numpy()
880
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
comfyanonymous's avatar
comfyanonymous committed
881
882
883
884
885
886
            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]))
887

888
            file = f"{filename}_{counter:05}_.png"
889
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
m957ymj75urz's avatar
m957ymj75urz committed
890
891
892
893
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
Gavroche CryptoRUSH's avatar
Gavroche CryptoRUSH committed
894
            })
895
            counter += 1
comfyanonymous's avatar
comfyanonymous committed
896

m957ymj75urz's avatar
m957ymj75urz committed
897
        return { "ui": { "images": results } }
comfyanonymous's avatar
comfyanonymous committed
898

pythongosssss's avatar
pythongosssss committed
899
900
class PreviewImage(SaveImage):
    def __init__(self):
901
        self.output_dir = folder_paths.get_temp_directory()
m957ymj75urz's avatar
m957ymj75urz committed
902
        self.type = "temp"
pythongosssss's avatar
pythongosssss committed
903
904
905

    @classmethod
    def INPUT_TYPES(s):
906
        return {"required":
pythongosssss's avatar
pythongosssss committed
907
908
909
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
910

911
912
913
class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
914
        input_dir = folder_paths.get_input_directory()
915
        return {"required":
916
                    {"image": (sorted(os.listdir(input_dir)), )},
917
                }
918
919

    CATEGORY = "image"
920

921
    RETURN_TYPES = ("IMAGE", "MASK")
922
923
    FUNCTION = "load_image"
    def load_image(self, image):
924
925
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
926
927
        i = Image.open(image_path)
        image = i.convert("RGB")
928
        image = np.array(image).astype(np.float32) / 255.0
929
        image = torch.from_numpy(image)[None,]
930
931
932
933
934
935
        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)
936

937
938
    @classmethod
    def IS_CHANGED(s, image):
939
940
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
941
942
943
944
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
945

946
947
948
class LoadImageMask:
    @classmethod
    def INPUT_TYPES(s):
949
        input_dir = folder_paths.get_input_directory()
950
        return {"required":
951
                    {"image": (sorted(os.listdir(input_dir)), ),
952
953
954
                    "channel": (["alpha", "red", "green", "blue"], ),}
                }

955
    CATEGORY = "mask"
956
957
958
959

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
960
961
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
962
        i = Image.open(image_path)
963
964
        if i.getbands() != ("R", "G", "B", "A"):
            i = i.convert("RGBA")
965
966
967
968
969
970
971
972
973
974
975
976
977
        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):
978
979
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
980
981
982
983
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
984

comfyanonymous's avatar
comfyanonymous committed
985
986
987
988
989
990
991
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,),
992
993
                              "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
994
995
996
997
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

998
    CATEGORY = "image/upscaling"
999

comfyanonymous's avatar
comfyanonymous committed
1000
1001
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
1002
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
1003
1004
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
1005

1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
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
1022
1023
1024
1025
1026
1027
1028
class ImagePadForOutpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
1029
1030
1031
1032
1033
                "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
1034
1035
1036
1037
1038
1039
1040
1041
            }
        }

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

    CATEGORY = "image"

1042
    def expand_image(self, image, left, top, right, bottom, feathering):
Guo Y.K's avatar
Guo Y.K committed
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
        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,
        )
1055

1056
1057
1058
1059
1060
        t = torch.zeros(
            (d2, d3),
            dtype=torch.float32
        )

1061
        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080

            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
1081

Guo Y.K's avatar
Guo Y.K committed
1082
1083
1084
        return (new_image, mask)


comfyanonymous's avatar
comfyanonymous committed
1085
1086
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
1087
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
1088
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
1089
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
1090
1091
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
1092
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
1093
1094
1095
1096
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
    "LatentUpscale": LatentUpscale,
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
1097
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
1098
    "LoadImage": LoadImage,
1099
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
1100
    "ImageScale": ImageScale,
1101
    "ImageInvert": ImageInvert,
Guo Y.K's avatar
Guo Y.K committed
1102
    "ImagePadForOutpaint": ImagePadForOutpaint,
comfyanonymous's avatar
comfyanonymous committed
1103
1104
    "ConditioningCombine": ConditioningCombine,
    "ConditioningSetArea": ConditioningSetArea,
comfyanonymous's avatar
comfyanonymous committed
1105
    "KSamplerAdvanced": KSamplerAdvanced,
1106
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
1107
    "LatentComposite": LatentComposite,
comfyanonymous's avatar
comfyanonymous committed
1108
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
1109
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
1110
    "LatentCrop": LatentCrop,
1111
    "LoraLoader": LoraLoader,
1112
    "CLIPLoader": CLIPLoader,
1113
    "CLIPVisionEncode": CLIPVisionEncode,
1114
    "StyleModelApply": StyleModelApply,
1115
    "unCLIPConditioning": unCLIPConditioning,
comfyanonymous's avatar
comfyanonymous committed
1116
1117
    "ControlNetApply": ControlNetApply,
    "ControlNetLoader": ControlNetLoader,
1118
    "DiffControlNetLoader": DiffControlNetLoader,
comfyanonymous's avatar
comfyanonymous committed
1119
1120
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
1121
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
1122
    "VAEEncodeTiled": VAEEncodeTiled,
1123
    "TomePatchModel": TomePatchModel,
1124
    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
1125
    "CheckpointLoader": CheckpointLoader,
sALTaccount's avatar
sALTaccount committed
1126
    "DiffusersLoader": DiffusersLoader,
comfyanonymous's avatar
comfyanonymous committed
1127
1128
}

City's avatar
City committed
1129
1130
1131
1132
1133
NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "KSampler": "KSampler",
    "KSamplerAdvanced": "KSampler (Advanced)",
    # Loaders
comfyanonymous's avatar
Rename.  
comfyanonymous committed
1134
1135
    "CheckpointLoader": "Load Checkpoint (With Config)",
    "CheckpointLoaderSimple": "Load Checkpoint",
City's avatar
City committed
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
    "VAELoader": "Load VAE",
    "LoraLoader": "Load LoRA",
    "CLIPLoader": "Load CLIP",
    "ControlNetLoader": "Load ControlNet Model",
    "DiffControlNetLoader": "Load ControlNet Model (diff)",
    "StyleModelLoader": "Load Style Model",
    "CLIPVisionLoader": "Load CLIP Vision",
    "UpscaleModelLoader": "Load Upscale Model",
    # Conditioning
    "CLIPVisionEncode": "CLIP Vision Encode",
    "StyleModelApply": "Apply Style Model",
    "CLIPTextEncode": "CLIP Text Encode (Prompt)",
    "CLIPSetLastLayer": "CLIP Set Last Layer",
    "ConditioningCombine": "Conditioning (Combine)",
    "ConditioningSetArea": "Conditioning (Set Area)",
    "ControlNetApply": "Apply ControlNet",
    # Latent
    "VAEEncodeForInpaint": "VAE Encode (for Inpainting)",
    "SetLatentNoiseMask": "Set Latent Noise Mask",
    "VAEDecode": "VAE Decode",
    "VAEEncode": "VAE Encode",
    "LatentRotate": "Rotate Latent",
    "LatentFlip": "Flip Latent",
    "LatentCrop": "Crop Latent",
    "EmptyLatentImage": "Empty Latent Image",
    "LatentUpscale": "Upscale Latent",
    "LatentComposite": "Latent Composite",
    # Image
    "SaveImage": "Save Image",
    "PreviewImage": "Preview Image",
    "LoadImage": "Load Image",
    "LoadImageMask": "Load Image (as Mask)",
    "ImageScale": "Upscale Image",
    "ImageUpscaleWithModel": "Upscale Image (using Model)",
    "ImageInvert": "Invert Image",
    "ImagePadForOutpaint": "Pad Image for Outpainting",
    # _for_testing
    "VAEDecodeTiled": "VAE Decode (Tiled)",
    "VAEEncodeTiled": "VAE Encode (Tiled)",
}

1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
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)
1192
1193
            if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
                NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
1194
1195
1196
1197
1198
1199
        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
1200
def load_custom_nodes():
1201
    CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
1202
    possible_modules = os.listdir(CUSTOM_NODE_PATH)
1203
    if "__pycache__" in possible_modules:
Hacker 17082006's avatar
.  
Hacker 17082006 committed
1204
        possible_modules.remove("__pycache__")
1205

Hacker 17082006's avatar
Hacker 17082006 committed
1206
    for possible_module in possible_modules:
1207
1208
        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
1209
        load_custom_node(module_path)
1210

1211
1212
def init_custom_nodes():
    load_custom_nodes()
1213
1214
    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"))
1215
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))