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

import os
import sys
import json
6
import hashlib
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

Jake D's avatar
Jake D committed
513
514
515
516
517
518
519
520
class Note:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"text": ("STRING", {"multiline": True})}}

    CATEGORY = "other"
    RETURN_TYPES = ()

comfyanonymous's avatar
comfyanonymous committed
521

comfyanonymous's avatar
comfyanonymous committed
522
523
class LatentUpscale:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
524
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
525
526
527
528

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
529
530
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
531
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
532
533
534
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

535
536
    CATEGORY = "latent"

537
    def upscale(self, samples, upscale_method, width, height, crop):
538
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
539
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
540
541
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
542
543
544
545
546
547
548
549
550
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
551
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
552
553

    def rotate(self, samples, rotation):
554
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
555
556
557
558
559
560
561
562
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

563
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
564
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
565
566
567
568
569
570
571
572
573
574

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
575
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
576
577

    def flip(self, samples, flip_method):
578
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
579
        if flip_method.startswith("x"):
580
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
581
        elif flip_method.startswith("y"):
582
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
583
584

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
585
586
587
588

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
589
590
591
592
593
594
        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
595
596
597
598
599
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
600
601
602
    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
        x =  x // 8
        y = y // 8
603
        feather = feather // 8
comfyanonymous's avatar
comfyanonymous committed
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
        samples_out = samples_to.copy()
        s = samples_to["samples"].clone()
        samples_to = samples_to["samples"]
        samples_from = samples_from["samples"]
        if feather == 0:
            s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
        else:
            samples_from = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
            mask = torch.ones_like(samples_from)
            for t in range(feather):
                if y != 0:
                    mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))

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

comfyanonymous's avatar
comfyanonymous committed
628
629
630
631
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
632
633
634
635
                              "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
636
637
638
639
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
640
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
641
642

    def crop(self, samples, width, height, x, y):
643
644
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
        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])
668
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
669
670
        return (s,)

671
672
673
674
675
676
677
678
679
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

680
    CATEGORY = "latent/inpaint"
681
682
683
684
685
686
687

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


688
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):
689
690
    latent_image = latent["samples"]
    noise_mask = None
691
    device = model_management.get_torch_device()
692

comfyanonymous's avatar
comfyanonymous committed
693
694
695
696
697
    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")

698
699
700
    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")
701
        noise_mask = noise_mask.round()
702
703
704
705
        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)

706
    real_model = None
707
708
709
    model_management.load_model_gpu(model)
    real_model = model.model

710
711
712
713
714
715
    noise = noise.to(device)
    latent_image = latent_image.to(device)

    positive_copy = []
    negative_copy = []

comfyanonymous's avatar
comfyanonymous committed
716
    control_nets = []
717
718
719
720
721
    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
722
723
        if 'control' in p[1]:
            control_nets += [p[1]['control']]
724
725
726
727
728
729
        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)
730
731
        if 'control' in n[1]:
            control_nets += [n[1]['control']]
732
733
        negative_copy += [[t] + n[1:]]

comfyanonymous's avatar
comfyanonymous committed
734
735
736
737
    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
738

739
    if sampler_name in comfy.samplers.KSampler.SAMPLERS:
740
        sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
741
742
743
744
    else:
        #other samplers
        pass

745
    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)
746
    samples = samples.cpu()
comfyanonymous's avatar
comfyanonymous committed
747
748
    for c in control_nets:
        c.cleanup()
comfyanonymous's avatar
comfyanonymous committed
749

750
751
752
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
753

comfyanonymous's avatar
comfyanonymous committed
754
755
756
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
757
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
                    {"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"

773
774
    CATEGORY = "sampling"

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

comfyanonymous's avatar
comfyanonymous committed
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
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
801

comfyanonymous's avatar
comfyanonymous committed
802
803
804
805
806
807
808
    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
809
        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
810
811
812

class SaveImage:
    def __init__(self):
813
        self.output_dir = folder_paths.get_output_directory()
m957ymj75urz's avatar
m957ymj75urz committed
814
        self.type = "output"
comfyanonymous's avatar
comfyanonymous committed
815
816
817
818

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
819
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
820
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
821
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
822
823
824
825
826
827
828
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

829
830
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
831
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
832
        def map_filename(filename):
833
            prefix_len = len(os.path.basename(filename_prefix))
834
835
836
837
838
839
            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
840

841
842
843
844
        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
845

846
        filename_prefix = compute_vars(filename_prefix)
comfyanonymous's avatar
comfyanonymous committed
847

m957ymj75urz's avatar
m957ymj75urz committed
848
849
850
        subfolder = os.path.dirname(os.path.normpath(filename_prefix))
        filename = os.path.basename(os.path.normpath(filename_prefix))

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

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

857
        try:
858
            counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
859
860
        except ValueError:
            counter = 1
861
        except FileNotFoundError:
862
            os.makedirs(full_output_folder, exist_ok=True)
863
            counter = 1
pythongosssss's avatar
pythongosssss committed
864

m957ymj75urz's avatar
m957ymj75urz committed
865
        results = list()
comfyanonymous's avatar
comfyanonymous committed
866
867
        for image in images:
            i = 255. * image.cpu().numpy()
868
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
comfyanonymous's avatar
comfyanonymous committed
869
870
871
872
873
874
            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]))
875

876
            file = f"{filename}_{counter:05}_.png"
877
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
m957ymj75urz's avatar
m957ymj75urz committed
878
879
880
881
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
Gavroche CryptoRUSH's avatar
Gavroche CryptoRUSH committed
882
            })
883
            counter += 1
comfyanonymous's avatar
comfyanonymous committed
884

m957ymj75urz's avatar
m957ymj75urz committed
885
        return { "ui": { "images": results } }
comfyanonymous's avatar
comfyanonymous committed
886

pythongosssss's avatar
pythongosssss committed
887
888
class PreviewImage(SaveImage):
    def __init__(self):
889
        self.output_dir = folder_paths.get_temp_directory()
m957ymj75urz's avatar
m957ymj75urz committed
890
        self.type = "temp"
pythongosssss's avatar
pythongosssss committed
891
892
893

    @classmethod
    def INPUT_TYPES(s):
894
        return {"required":
pythongosssss's avatar
pythongosssss committed
895
896
897
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
898

899
900
901
class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
902
        input_dir = folder_paths.get_input_directory()
903
        return {"required":
904
                    {"image": (sorted(os.listdir(input_dir)), )},
905
                }
906
907

    CATEGORY = "image"
908

909
    RETURN_TYPES = ("IMAGE", "MASK")
910
911
    FUNCTION = "load_image"
    def load_image(self, image):
912
913
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
914
915
        i = Image.open(image_path)
        image = i.convert("RGB")
916
        image = np.array(image).astype(np.float32) / 255.0
917
        image = torch.from_numpy(image)[None,]
918
919
920
921
922
923
        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)
924

925
926
    @classmethod
    def IS_CHANGED(s, image):
927
928
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
929
930
931
932
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
933

934
935
936
class LoadImageMask:
    @classmethod
    def INPUT_TYPES(s):
937
        input_dir = folder_paths.get_input_directory()
938
        return {"required":
939
                    {"image": (sorted(os.listdir(input_dir)), ),
940
941
942
                    "channel": (["alpha", "red", "green", "blue"], ),}
                }

943
    CATEGORY = "mask"
944
945
946
947

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
948
949
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
950
        i = Image.open(image_path)
951
952
        if i.getbands() != ("R", "G", "B", "A"):
            i = i.convert("RGBA")
953
954
955
956
957
958
959
960
961
962
963
964
965
        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):
966
967
        input_dir = folder_paths.get_input_directory()
        image_path = os.path.join(input_dir, image)
968
969
970
971
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
972

comfyanonymous's avatar
comfyanonymous committed
973
974
975
976
977
978
979
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,),
980
981
                              "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
982
983
984
985
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

986
    CATEGORY = "image/upscaling"
987

comfyanonymous's avatar
comfyanonymous committed
988
989
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
990
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
991
992
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
993

994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
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
1010
1011
1012
1013
1014
1015
1016
class ImagePadForOutpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
1017
1018
1019
1020
1021
                "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
1022
1023
1024
1025
1026
1027
1028
1029
            }
        }

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

    CATEGORY = "image"

1030
    def expand_image(self, image, left, top, right, bottom, feathering):
Guo Y.K's avatar
Guo Y.K committed
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
        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,
        )
1043

1044
1045
1046
1047
1048
        t = torch.zeros(
            (d2, d3),
            dtype=torch.float32
        )

1049
        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068

            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
1069

Guo Y.K's avatar
Guo Y.K committed
1070
1071
1072
        return (new_image, mask)


comfyanonymous's avatar
comfyanonymous committed
1073
1074
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
1075
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
1076
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
1077
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
1078
1079
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
1080
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
1081
1082
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
Jake D's avatar
Jake D committed
1083
    "Note": Note,
comfyanonymous's avatar
comfyanonymous committed
1084
1085
    "LatentUpscale": LatentUpscale,
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
1086
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
1087
    "LoadImage": LoadImage,
1088
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
1089
    "ImageScale": ImageScale,
1090
    "ImageInvert": ImageInvert,
Guo Y.K's avatar
Guo Y.K committed
1091
    "ImagePadForOutpaint": ImagePadForOutpaint,
comfyanonymous's avatar
comfyanonymous committed
1092
1093
    "ConditioningCombine": ConditioningCombine,
    "ConditioningSetArea": ConditioningSetArea,
comfyanonymous's avatar
comfyanonymous committed
1094
    "KSamplerAdvanced": KSamplerAdvanced,
1095
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
1096
    "LatentComposite": LatentComposite,
comfyanonymous's avatar
comfyanonymous committed
1097
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
1098
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
1099
    "LatentCrop": LatentCrop,
1100
    "LoraLoader": LoraLoader,
1101
    "CLIPLoader": CLIPLoader,
1102
    "CLIPVisionEncode": CLIPVisionEncode,
1103
    "StyleModelApply": StyleModelApply,
1104
    "unCLIPConditioning": unCLIPConditioning,
comfyanonymous's avatar
comfyanonymous committed
1105
1106
    "ControlNetApply": ControlNetApply,
    "ControlNetLoader": ControlNetLoader,
1107
    "DiffControlNetLoader": DiffControlNetLoader,
comfyanonymous's avatar
comfyanonymous committed
1108
1109
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
1110
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
1111
    "VAEEncodeTiled": VAEEncodeTiled,
1112
    "TomePatchModel": TomePatchModel,
1113
    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
1114
    "CheckpointLoader": CheckpointLoader,
sALTaccount's avatar
sALTaccount committed
1115
    "DiffusersLoader": DiffusersLoader,
comfyanonymous's avatar
comfyanonymous committed
1116
1117
}

City's avatar
City committed
1118
1119
1120
1121
1122
NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "KSampler": "KSampler",
    "KSamplerAdvanced": "KSampler (Advanced)",
    # Loaders
comfyanonymous's avatar
Rename.  
comfyanonymous committed
1123
1124
    "CheckpointLoader": "Load Checkpoint (With Config)",
    "CheckpointLoaderSimple": "Load Checkpoint",
City's avatar
City committed
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
    "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",
Jake D's avatar
Jake D committed
1150
    "Note": "Note",
City's avatar
City committed
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
    "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)",
}

1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
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)
1182
1183
            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)
1184
1185
1186
1187
1188
1189
        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
1190
def load_custom_nodes():
1191
    CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
1192
    possible_modules = os.listdir(CUSTOM_NODE_PATH)
1193
    if "__pycache__" in possible_modules:
Hacker 17082006's avatar
.  
Hacker 17082006 committed
1194
        possible_modules.remove("__pycache__")
1195

Hacker 17082006's avatar
Hacker 17082006 committed
1196
    for possible_module in possible_modules:
1197
1198
        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
1199
        load_custom_node(module_path)
1200

1201
1202
def init_custom_nodes():
    load_custom_nodes()
1203
1204
    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"))
1205
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))