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

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

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

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


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

27
import comfy.clip_vision
28

29
import comfy.model_management
30
31
from comfy.cli_args import args

32
import importlib
comfyanonymous's avatar
comfyanonymous committed
33

34
import folder_paths
35
import latent_preview
space-nuko's avatar
space-nuko committed
36

37
def before_node_execution():
38
    comfy.model_management.throw_exception_if_processing_interrupted()
39

40
def interrupt_processing(value=True):
41
    comfy.model_management.interrupt_current_processing(value)
42

43
44
MAX_RESOLUTION=8192

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

52
53
    CATEGORY = "conditioning"

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

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

66
67
    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
68
69
70
    def combine(self, conditioning_1, conditioning_2):
        return (conditioning_1 + conditioning_2, )

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

    CATEGORY = "conditioning"

comfyanonymous's avatar
comfyanonymous committed
82
    def addWeighted(self, conditioning_to, conditioning_from, conditioning_to_strength):
FizzleDorf's avatar
FizzleDorf committed
83
        out = []
comfyanonymous's avatar
comfyanonymous committed
84
85
86
87
88

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

        cond_from = conditioning_from[0][0]
89
        pooled_output_from = conditioning_from[0][1].get("pooled_output", None)
comfyanonymous's avatar
comfyanonymous committed
90
91
92

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

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

            n = [tw, t_to]
FizzleDorf's avatar
FizzleDorf committed
106
107
108
            out.append(n)
        return (out, )

109
110
111
112
113
114
115
116
117
118
class ConditioningConcat:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "conditioning_to": ("CONDITIONING",),
            "conditioning_from": ("CONDITIONING",),
            }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "concat"

119
    CATEGORY = "conditioning"
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

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

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

        cond_from = conditioning_from[0][0]

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

        return (out, )

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

150
151
    CATEGORY = "conditioning"

152
    def append(self, conditioning, width, height, x, y, strength):
comfyanonymous's avatar
comfyanonymous committed
153
154
155
156
157
        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
158
            n[1]['set_area_to_bounds'] = False
comfyanonymous's avatar
comfyanonymous committed
159
            c.append(n)
comfyanonymous's avatar
comfyanonymous committed
160
        return (c, )
comfyanonymous's avatar
comfyanonymous committed
161

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

    CATEGORY = "conditioning"

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

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

    CATEGORY = "advanced/conditioning"

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

210
211
212
213
class ConditioningSetTimestepRange:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"conditioning": ("CONDITIONING", ),
214
215
                             "start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                             "end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
216
217
218
219
220
221
222
223
224
225
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "set_range"

    CATEGORY = "advanced/conditioning"

    def set_range(self, conditioning, start, end):
        c = []
        for t in conditioning:
            d = t[1].copy()
226
227
            d['start_percent'] = 1.0 - start
            d['end_percent'] = 1.0 - end
228
229
230
231
            n = [t[0], d]
            c.append(n)
        return (c, )

comfyanonymous's avatar
comfyanonymous committed
232
233
234
235
236
237
238
class VAEDecode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"

239
240
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
241
    def decode(self, vae, samples):
242
        return (vae.decode(samples["samples"]), )
comfyanonymous's avatar
comfyanonymous committed
243

244
245
246
247
248
249
250
251
252
253
254
255
class VAEDecodeTiled:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"

    CATEGORY = "_for_testing"

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

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

263
264
    CATEGORY = "latent"

265
266
267
268
    @staticmethod
    def vae_encode_crop_pixels(pixels):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
269
        if pixels.shape[1] != x or pixels.shape[2] != y:
270
271
272
273
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
        return pixels
274

275
276
277
    def encode(self, vae, pixels):
        pixels = self.vae_encode_crop_pixels(pixels)
        t = vae.encode(pixels[:,:,:,:3])
278
        return ({"samples":t}, )
comfyanonymous's avatar
comfyanonymous committed
279

comfyanonymous's avatar
comfyanonymous committed
280
281
282
283
284
285
286
287
288
289
class VAEEncodeTiled:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "_for_testing"

    def encode(self, vae, pixels):
290
        pixels = VAEEncode.vae_encode_crop_pixels(pixels)
comfyanonymous's avatar
comfyanonymous committed
291
292
        t = vae.encode_tiled(pixels[:,:,:,:3])
        return ({"samples":t}, )
293

294
295
296
class VAEEncodeForInpaint:
    @classmethod
    def INPUT_TYPES(s):
297
        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", ), "mask": ("MASK", ), "grow_mask_by": ("INT", {"default": 6, "min": 0, "max": 64, "step": 1}),}}
298
299
300
301
302
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"

    CATEGORY = "latent/inpaint"

303
    def encode(self, vae, pixels, mask, grow_mask_by=6):
304
305
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
306
        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
307

308
        pixels = pixels.clone()
309
        if pixels.shape[1] != x or pixels.shape[2] != y:
310
311
312
313
            x_offset = (pixels.shape[1] % 8) // 2
            y_offset = (pixels.shape[2] % 8) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
314

315
        #grow mask by a few pixels to keep things seamless in latent space
316
317
318
319
320
321
322
323
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

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

324
        m = (1.0 - mask.round()).squeeze(1)
325
326
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
327
            pixels[:,:,:,i] *= m
328
329
330
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

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

Dr.Lt.Data's avatar
Dr.Lt.Data committed
333
334
class SaveLatent:
    def __init__(self):
335
        self.output_dir = folder_paths.get_output_directory()
Dr.Lt.Data's avatar
Dr.Lt.Data committed
336
337
338
339

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT", ),
340
                              "filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
Dr.Lt.Data's avatar
Dr.Lt.Data committed
341
342
343
344
345
346
347
348
349
350
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
    RETURN_TYPES = ()
    FUNCTION = "save"

    OUTPUT_NODE = True

    CATEGORY = "_for_testing"

    def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
351
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
Dr.Lt.Data's avatar
Dr.Lt.Data committed
352
353
354
355
356
357

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

358
359
360
361
362
363
        metadata = None
        if not args.disable_metadata:
            metadata = {"prompt": prompt_info}
            if extra_pnginfo is not None:
                for x in extra_pnginfo:
                    metadata[x] = json.dumps(extra_pnginfo[x])
Dr.Lt.Data's avatar
Dr.Lt.Data committed
364
365

        file = f"{filename}_{counter:05}_.latent"
366
367
368
369
370
371
372
373

        results = list()
        results.append({
            "filename": file,
            "subfolder": subfolder,
            "type": "output"
        })

Dr.Lt.Data's avatar
Dr.Lt.Data committed
374
375
        file = os.path.join(full_output_folder, file)

376
377
        output = {}
        output["latent_tensor"] = samples["samples"]
378
        output["latent_format_version_0"] = torch.tensor([])
379

380
        comfy.utils.save_torch_file(output, file, metadata=metadata)
381
        return { "ui": { "latents": results } }
Dr.Lt.Data's avatar
Dr.Lt.Data committed
382
383
384
385
386


class LoadLatent:
    @classmethod
    def INPUT_TYPES(s):
387
388
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".latent")]
Dr.Lt.Data's avatar
Dr.Lt.Data committed
389
390
391
392
393
394
395
396
        return {"required": {"latent": [sorted(files), ]}, }

    CATEGORY = "_for_testing"

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

    def load(self, latent):
397
398
        latent_path = folder_paths.get_annotated_filepath(latent)
        latent = safetensors.torch.load_file(latent_path, device="cpu")
399
400
401
402
        multiplier = 1.0
        if "latent_format_version_0" not in latent:
            multiplier = 1.0 / 0.18215
        samples = {"samples": latent["latent_tensor"].float() * multiplier}
403
        return (samples, )
Dr.Lt.Data's avatar
Dr.Lt.Data committed
404

405
406
407
408
409
410
411
412
413
414
415
416
417
418
    @classmethod
    def IS_CHANGED(s, latent):
        image_path = folder_paths.get_annotated_filepath(latent)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()

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

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

comfyanonymous's avatar
comfyanonymous committed
420
421
422
class CheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
423
424
        return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ),
                              "ckpt_name": (folder_paths.get_filename_list("checkpoints"), )}}
comfyanonymous's avatar
comfyanonymous committed
425
426
427
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

428
    CATEGORY = "advanced/loaders"
429

comfyanonymous's avatar
comfyanonymous committed
430
    def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
431
432
        config_path = folder_paths.get_full_path("configs", config_name)
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
433
        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
434

435
436
437
class CheckpointLoaderSimple:
    @classmethod
    def INPUT_TYPES(s):
438
        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
439
440
441
442
                             }}
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

443
    CATEGORY = "loaders"
444

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

sALTaccount's avatar
sALTaccount committed
450
451
452
class DiffusersLoader:
    @classmethod
    def INPUT_TYPES(cls):
453
        paths = []
sALTaccount's avatar
sALTaccount committed
454
        for search_path in folder_paths.get_folder_paths("diffusers"):
455
            if os.path.exists(search_path):
456
457
458
459
                for root, subdir, files in os.walk(search_path, followlinks=True):
                    if "model_index.json" in files:
                        paths.append(os.path.relpath(root, start=search_path))

460
        return {"required": {"model_path": (paths,), }}
sALTaccount's avatar
sALTaccount committed
461
462
463
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

464
    CATEGORY = "advanced/loaders/deprecated"
sALTaccount's avatar
sALTaccount committed
465
466

    def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
sALTaccount's avatar
sALTaccount committed
467
468
        for search_path in folder_paths.get_folder_paths("diffusers"):
            if os.path.exists(search_path):
469
470
471
                path = os.path.join(search_path, model_path)
                if os.path.exists(path):
                    model_path = path
sALTaccount's avatar
sALTaccount committed
472
                    break
473

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


477
478
479
480
481
482
483
484
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"

485
    CATEGORY = "loaders"
486
487
488
489
490
491

    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
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
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,)

508
class LoraLoader:
509
510
511
    def __init__(self):
        self.loaded_lora = None

512
513
514
515
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
                              "clip": ("CLIP", ),
516
                              "lora_name": (folder_paths.get_filename_list("loras"), ),
517
518
                              "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}),
519
520
521
522
523
524
525
                              }}
    RETURN_TYPES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"

    CATEGORY = "loaders"

    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
526
527
528
        if strength_model == 0 and strength_clip == 0:
            return (model, clip)

529
        lora_path = folder_paths.get_full_path("loras", lora_name)
530
531
532
533
534
        lora = None
        if self.loaded_lora is not None:
            if self.loaded_lora[0] == lora_path:
                lora = self.loaded_lora[1]
            else:
535
536
537
                temp = self.loaded_lora
                self.loaded_lora = None
                del temp
538
539
540
541
542
543

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

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

comfyanonymous's avatar
comfyanonymous committed
546
547
548
class VAELoader:
    @classmethod
    def INPUT_TYPES(s):
549
        return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}}
comfyanonymous's avatar
comfyanonymous committed
550
551
552
    RETURN_TYPES = ("VAE",)
    FUNCTION = "load_vae"

553
554
    CATEGORY = "loaders"

comfyanonymous's avatar
comfyanonymous committed
555
556
    #TODO: scale factor?
    def load_vae(self, vae_name):
557
        vae_path = folder_paths.get_full_path("vae", vae_name)
comfyanonymous's avatar
comfyanonymous committed
558
559
560
        vae = comfy.sd.VAE(ckpt_path=vae_path)
        return (vae,)

comfyanonymous's avatar
comfyanonymous committed
561
562
563
class ControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
564
        return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
comfyanonymous's avatar
comfyanonymous committed
565
566
567
568
569
570
571

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

    CATEGORY = "loaders"

    def load_controlnet(self, control_net_name):
572
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
573
        controlnet = comfy.controlnet.load_controlnet(controlnet_path)
comfyanonymous's avatar
comfyanonymous committed
574
575
        return (controlnet,)

576
577
578
579
class DiffControlNetLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model": ("MODEL",),
580
                              "control_net_name": (folder_paths.get_filename_list("controlnet"), )}}
581
582
583
584
585
586
587

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

    CATEGORY = "loaders"

    def load_controlnet(self, model, control_net_name):
588
        controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
589
        controlnet = comfy.controlnet.load_controlnet(controlnet_path, model)
590
591
        return (controlnet,)

comfyanonymous's avatar
comfyanonymous committed
592
593
594
595

class ControlNetApply:
    @classmethod
    def INPUT_TYPES(s):
596
597
598
599
600
        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
601
602
603
604
605
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

606
    def apply_controlnet(self, conditioning, control_net, image, strength):
607
608
609
        if strength == 0:
            return (conditioning, )

comfyanonymous's avatar
comfyanonymous committed
610
611
612
613
        c = []
        control_hint = image.movedim(-1,1)
        for t in conditioning:
            n = [t[0], t[1].copy()]
comfyanonymous's avatar
comfyanonymous committed
614
615
616
617
            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
618
            n[1]['control_apply_to_uncond'] = True
comfyanonymous's avatar
comfyanonymous committed
619
620
621
            c.append(n)
        return (c, )

622
623
624
625
626
627
628
629
630

class ControlNetApplyAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"positive": ("CONDITIONING", ),
                             "negative": ("CONDITIONING", ),
                             "control_net": ("CONTROL_NET", ),
                             "image": ("IMAGE", ),
                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
631
632
                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
633
634
635
636
637
638
639
640
                             }}

    RETURN_TYPES = ("CONDITIONING","CONDITIONING")
    RETURN_NAMES = ("positive", "negative")
    FUNCTION = "apply_controlnet"

    CATEGORY = "conditioning"

641
    def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent):
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        if strength == 0:
            return (positive, negative)

        control_hint = image.movedim(-1,1)
        cnets = {}

        out = []
        for conditioning in [positive, negative]:
            c = []
            for t in conditioning:
                d = t[1].copy()

                prev_cnet = d.get('control', None)
                if prev_cnet in cnets:
                    c_net = cnets[prev_cnet]
                else:
658
                    c_net = control_net.copy().set_cond_hint(control_hint, strength, (1.0 - start_percent, 1.0 - end_percent))
659
660
661
662
663
664
665
666
667
668
669
                    c_net.set_previous_controlnet(prev_cnet)
                    cnets[prev_cnet] = c_net

                d['control'] = c_net
                d['control_apply_to_uncond'] = False
                n = [t[0], d]
                c.append(n)
            out.append(c)
        return (out[0], out[1])


670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
class UNETLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ),
                             }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"

    CATEGORY = "advanced/loaders"

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

685
686
687
class CLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
688
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
689
690
691
692
                             }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "load_clip"

693
    CATEGORY = "advanced/loaders"
694

695
    def load_clip(self, clip_name):
696
        clip_path = folder_paths.get_full_path("clip", clip_name)
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
        clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"))
        return (clip,)

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

    CATEGORY = "advanced/loaders"

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

716
717
718
class CLIPVisionLoader:
    @classmethod
    def INPUT_TYPES(s):
719
        return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ),
720
721
722
723
724
725
726
                             }}
    RETURN_TYPES = ("CLIP_VISION",)
    FUNCTION = "load_clip"

    CATEGORY = "loaders"

    def load_clip(self, clip_name):
727
        clip_path = folder_paths.get_full_path("clip_vision", clip_name)
728
        clip_vision = comfy.clip_vision.load(clip_path)
729
730
731
732
733
734
735
736
        return (clip_vision,)

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

740
    CATEGORY = "conditioning"
741
742
743
744
745
746
747
748

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

class StyleModelLoader:
    @classmethod
    def INPUT_TYPES(s):
749
        return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}}
750
751
752
753
754
755
756

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

    CATEGORY = "loaders"

    def load_style_model(self, style_model_name):
757
        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
758
759
760
761
762
763
764
        style_model = comfy.sd.load_style_model(style_model_path)
        return (style_model,)


class StyleModelApply:
    @classmethod
    def INPUT_TYPES(s):
765
766
767
        return {"required": {"conditioning": ("CONDITIONING", ),
                             "style_model": ("STYLE_MODEL", ),
                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
768
769
770
771
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_stylemodel"

comfyanonymous's avatar
comfyanonymous committed
772
    CATEGORY = "conditioning/style_model"
773

774
    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
775
        cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
776
        c = []
777
778
        for t in conditioning:
            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
779
780
781
            c.append(n)
        return (c, )

782
783
784
785
786
787
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}),
788
                             "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
789
790
791
792
                             }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "apply_adm"

793
    CATEGORY = "conditioning"
794

795
    def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation):
796
797
798
        if strength == 0:
            return (conditioning, )

799
800
801
        c = []
        for t in conditioning:
            o = t[1].copy()
802
803
804
            x = {"clip_vision_output": clip_vision_output, "strength": strength, "noise_augmentation": noise_augmentation}
            if "unclip_conditioning" in o:
                o["unclip_conditioning"] = o["unclip_conditioning"][:] + [x]
805
            else:
806
                o["unclip_conditioning"] = [x]
807
808
809
810
            n = [t[0], o]
            c.append(n)
        return (c, )

811
812
813
814
815
816
817
818
class GLIGENLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "gligen_name": (folder_paths.get_filename_list("gligen"), )}}

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

comfyanonymous's avatar
comfyanonymous committed
819
    CATEGORY = "loaders"
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840

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

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

comfyanonymous's avatar
comfyanonymous committed
841
    CATEGORY = "conditioning/gligen"
842
843
844
845
846
847
848
849
850
851
852
853
854
855

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

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

comfyanonymous's avatar
comfyanonymous committed
857
858
859
860
861
862
class EmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
863
864
        return {"required": { "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
comfyanonymous's avatar
comfyanonymous committed
865
866
867
868
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

869
870
    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
871
872
    def generate(self, width, height, batch_size=1):
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
873
        return ({"samples":latent}, )
comfyanonymous's avatar
comfyanonymous committed
874

comfyanonymous's avatar
comfyanonymous committed
875

876
877
878
879
880
class LatentFromBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "batch_index": ("INT", {"default": 0, "min": 0, "max": 63}),
881
                              "length": ("INT", {"default": 1, "min": 1, "max": 64}),
882
883
                              }}
    RETURN_TYPES = ("LATENT",)
884
    FUNCTION = "frombatch"
885

886
    CATEGORY = "latent/batch"
887

888
    def frombatch(self, samples, batch_index, length):
889
890
891
        s = samples.copy()
        s_in = samples["samples"]
        batch_index = min(s_in.shape[0] - 1, batch_index)
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
        length = min(s_in.shape[0] - batch_index, length)
        s["samples"] = s_in[batch_index:batch_index + length].clone()
        if "noise_mask" in samples:
            masks = samples["noise_mask"]
            if masks.shape[0] == 1:
                s["noise_mask"] = masks.clone()
            else:
                if masks.shape[0] < s_in.shape[0]:
                    masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
                s["noise_mask"] = masks[batch_index:batch_index + length].clone()
        if "batch_index" not in s:
            s["batch_index"] = [x for x in range(batch_index, batch_index+length)]
        else:
            s["batch_index"] = samples["batch_index"][batch_index:batch_index + length]
        return (s,)
    
class RepeatLatentBatch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "amount": ("INT", {"default": 1, "min": 1, "max": 64}),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "repeat"

    CATEGORY = "latent/batch"

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

comfyanonymous's avatar
comfyanonymous committed
934
class LatentUpscale:
comfyanonymous's avatar
comfyanonymous committed
935
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
936
    crop_methods = ["disabled", "center"]
comfyanonymous's avatar
comfyanonymous committed
937
938
939
940

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
941
942
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
943
                              "crop": (s.crop_methods,)}}
comfyanonymous's avatar
comfyanonymous committed
944
945
946
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"

947
948
    CATEGORY = "latent"

949
    def upscale(self, samples, upscale_method, width, height, crop):
950
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
951
        s["samples"] = comfy.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
952
953
        return (s,)

comfyanonymous's avatar
comfyanonymous committed
954
class LatentUpscaleBy:
comfyanonymous's avatar
comfyanonymous committed
955
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
comfyanonymous's avatar
comfyanonymous committed
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972

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

    CATEGORY = "latent"

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

comfyanonymous's avatar
comfyanonymous committed
973
974
975
976
977
978
979
980
981
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
982
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
983
984

    def rotate(self, samples, rotation):
985
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
986
987
988
989
990
991
992
993
        rotate_by = 0
        if rotation.startswith("90"):
            rotate_by = 1
        elif rotation.startswith("180"):
            rotate_by = 2
        elif rotation.startswith("270"):
            rotate_by = 3

994
        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
comfyanonymous's avatar
comfyanonymous committed
995
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
996
997
998
999
1000
1001
1002
1003
1004
1005

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
1006
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
1007
1008

    def flip(self, samples, flip_method):
1009
        s = samples.copy()
comfyanonymous's avatar
comfyanonymous committed
1010
        if flip_method.startswith("x"):
1011
            s["samples"] = torch.flip(samples["samples"], dims=[2])
comfyanonymous's avatar
comfyanonymous committed
1012
        elif flip_method.startswith("y"):
1013
            s["samples"] = torch.flip(samples["samples"], dims=[3])
comfyanonymous's avatar
comfyanonymous committed
1014
1015

        return (s,)
comfyanonymous's avatar
comfyanonymous committed
1016
1017
1018
1019

class LatentComposite:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
1020
1021
1022
1023
1024
1025
        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
1026
1027
1028
1029
1030
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "composite"

    CATEGORY = "latent"

comfyanonymous's avatar
comfyanonymous committed
1031
1032
1033
    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
        x =  x // 8
        y = y // 8
1034
        feather = feather // 8
comfyanonymous's avatar
comfyanonymous committed
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
        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
1058

1059
1060
1061
1062
class LatentBlend:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
1063
1064
            "samples1": ("LATENT",),
            "samples2": ("LATENT",),
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
            "blend_factor": ("FLOAT", {
                "default": 0.5,
                "min": 0,
                "max": 1,
                "step": 0.01
            }),
        }}

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

    CATEGORY = "_for_testing"

1078
    def blend(self, samples1, samples2, blend_factor:float, blend_mode: str="normal"):
1079

1080
1081
1082
        samples_out = samples1.copy()
        samples1 = samples1["samples"]
        samples2 = samples2["samples"]
1083

1084
1085
1086
1087
        if samples1.shape != samples2.shape:
            samples2.permute(0, 3, 1, 2)
            samples2 = comfy.utils.common_upscale(samples2, samples1.shape[3], samples1.shape[2], 'bicubic', crop='center')
            samples2.permute(0, 2, 3, 1)
1088

1089
1090
        samples_blended = self.blend_mode(samples1, samples2, blend_mode)
        samples_blended = samples1 * blend_factor + samples_blended * (1 - blend_factor)
1091
1092
1093
1094
1095
1096
1097
1098
1099
        samples_out["samples"] = samples_blended
        return (samples_out,)

    def blend_mode(self, img1, img2, mode):
        if mode == "normal":
            return img2
        else:
            raise ValueError(f"Unsupported blend mode: {mode}")

comfyanonymous's avatar
comfyanonymous committed
1100
1101
1102
1103
class LatentCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
1104
1105
                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
1106
1107
                              "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
1108
1109
1110
1111
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "crop"

comfyanonymous's avatar
comfyanonymous committed
1112
    CATEGORY = "latent/transform"
comfyanonymous's avatar
comfyanonymous committed
1113
1114

    def crop(self, samples, width, height, x, y):
1115
1116
        s = samples.copy()
        samples = samples['samples']
comfyanonymous's avatar
comfyanonymous committed
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
        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
1130
        s['samples'] = samples[:,:,y:to_y, x:to_x]
comfyanonymous's avatar
comfyanonymous committed
1131
1132
        return (s,)

1133
1134
1135
1136
1137
1138
1139
1140
1141
class SetLatentNoiseMask:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "set_mask"

1142
    CATEGORY = "latent/inpaint"
1143
1144
1145

    def set_mask(self, samples, mask):
        s = samples.copy()
1146
        s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
1147
1148
        return (s,)

space-nuko's avatar
space-nuko committed
1149

space-nuko's avatar
space-nuko committed
1150
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):
1151
    device = comfy.model_management.get_torch_device()
1152
    latent_image = latent["samples"]
1153

comfyanonymous's avatar
comfyanonymous committed
1154
1155
1156
    if disable_noise:
        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
    else:
1157
1158
        batch_inds = latent["batch_index"] if "batch_index" in latent else None
        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
comfyanonymous's avatar
comfyanonymous committed
1159

1160
    noise_mask = None
1161
    if "noise_mask" in latent:
1162
        noise_mask = latent["noise_mask"]
comfyanonymous's avatar
comfyanonymous committed
1163

space-nuko's avatar
space-nuko committed
1164
1165
1166
1167
    preview_format = "JPEG"
    if preview_format not in ["JPEG", "PNG"]:
        preview_format = "JPEG"

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

1170
    pbar = comfy.utils.ProgressBar(steps)
1171
    def callback(step, x0, x, total_steps):
space-nuko's avatar
space-nuko committed
1172
        preview_bytes = None
1173
        if previewer:
1174
            preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
space-nuko's avatar
space-nuko committed
1175
        pbar.update_absolute(step + 1, total_steps, preview_bytes)
1176

1177
1178
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
1179
                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, seed=seed)
1180
1181
1182
    out = latent.copy()
    out["samples"] = samples
    return (out, )
comfyanonymous's avatar
comfyanonymous committed
1183

comfyanonymous's avatar
comfyanonymous committed
1184
1185
1186
class KSampler:
    @classmethod
    def INPUT_TYPES(s):
comfyanonymous's avatar
comfyanonymous committed
1187
        return {"required":
comfyanonymous's avatar
comfyanonymous committed
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
                    {"model": ("MODEL",),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "latent_image": ("LATENT", ),
                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
space-nuko's avatar
space-nuko committed
1198
1199
                     }
                }
comfyanonymous's avatar
comfyanonymous committed
1200
1201
1202
1203

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

1204
1205
    CATEGORY = "sampling"

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

comfyanonymous's avatar
comfyanonymous committed
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
class KSamplerAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL",),
                    "add_noise": (["enable", "disable"], ),
                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "latent_image": ("LATENT", ),
                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
                    "return_with_leftover_noise": (["disable", "enable"], ),
space-nuko's avatar
space-nuko committed
1226
1227
                     }
                }
comfyanonymous's avatar
comfyanonymous committed
1228
1229
1230
1231
1232

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

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

space-nuko's avatar
space-nuko committed
1234
    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
comfyanonymous's avatar
comfyanonymous committed
1235
1236
1237
1238
1239
1240
        force_full_denoise = True
        if return_with_leftover_noise == "enable":
            force_full_denoise = False
        disable_noise = False
        if add_noise == "disable":
            disable_noise = True
space-nuko's avatar
space-nuko committed
1241
        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
1242
1243
1244

class SaveImage:
    def __init__(self):
1245
        self.output_dir = folder_paths.get_output_directory()
m957ymj75urz's avatar
m957ymj75urz committed
1246
        self.type = "output"
1247
        self.prefix_append = ""
comfyanonymous's avatar
comfyanonymous committed
1248
1249
1250
1251

    @classmethod
    def INPUT_TYPES(s):
        return {"required": 
1252
                    {"images": ("IMAGE", ),
pythongosssss's avatar
tidy  
pythongosssss committed
1253
                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
pythongosssss's avatar
pythongosssss committed
1254
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
comfyanonymous's avatar
comfyanonymous committed
1255
1256
1257
1258
1259
1260
1261
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

1262
1263
    CATEGORY = "image"

pythongosssss's avatar
tidy  
pythongosssss committed
1264
    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
1265
        filename_prefix += self.prefix_append
1266
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
m957ymj75urz's avatar
m957ymj75urz committed
1267
        results = list()
comfyanonymous's avatar
comfyanonymous committed
1268
1269
        for image in images:
            i = 255. * image.cpu().numpy()
1270
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
1271
1272
1273
1274
1275
1276
1277
1278
            metadata = None
            if not args.disable_metadata:
                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]))
1279

1280
            file = f"{filename}_{counter:05}_.png"
1281
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
m957ymj75urz's avatar
m957ymj75urz committed
1282
1283
1284
1285
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
Gavroche CryptoRUSH's avatar
Gavroche CryptoRUSH committed
1286
            })
1287
            counter += 1
comfyanonymous's avatar
comfyanonymous committed
1288

m957ymj75urz's avatar
m957ymj75urz committed
1289
        return { "ui": { "images": results } }
comfyanonymous's avatar
comfyanonymous committed
1290

pythongosssss's avatar
pythongosssss committed
1291
1292
class PreviewImage(SaveImage):
    def __init__(self):
1293
        self.output_dir = folder_paths.get_temp_directory()
m957ymj75urz's avatar
m957ymj75urz committed
1294
        self.type = "temp"
1295
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
pythongosssss's avatar
pythongosssss committed
1296
1297
1298

    @classmethod
    def INPUT_TYPES(s):
1299
        return {"required":
pythongosssss's avatar
pythongosssss committed
1300
1301
1302
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }
1303

1304
1305
1306
class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
1307
        input_dir = folder_paths.get_input_directory()
comfyanonymous's avatar
comfyanonymous committed
1308
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
1309
        return {"required":
1310
                    {"image": (sorted(files), {"image_upload": True})},
1311
                }
1312
1313

    CATEGORY = "image"
1314

1315
    RETURN_TYPES = ("IMAGE", "MASK")
1316
1317
    FUNCTION = "load_image"
    def load_image(self, image):
1318
        image_path = folder_paths.get_annotated_filepath(image)
1319
        i = Image.open(image_path)
1320
        i = ImageOps.exif_transpose(i)
1321
        image = i.convert("RGB")
1322
        image = np.array(image).astype(np.float32) / 255.0
1323
        image = torch.from_numpy(image)[None,]
1324
1325
1326
1327
1328
1329
        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)
1330

1331
1332
    @classmethod
    def IS_CHANGED(s, image):
1333
        image_path = folder_paths.get_annotated_filepath(image)
1334
1335
1336
1337
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
1338

1339
1340
1341
1342
1343
1344
1345
    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        return True

1346
class LoadImageMask:
1347
    _color_channels = ["alpha", "red", "green", "blue"]
1348
1349
    @classmethod
    def INPUT_TYPES(s):
1350
        input_dir = folder_paths.get_input_directory()
comfyanonymous's avatar
comfyanonymous committed
1351
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
1352
        return {"required":
1353
                    {"image": (sorted(files), {"image_upload": True}),
1354
                     "channel": (s._color_channels, ), }
1355
1356
                }

1357
    CATEGORY = "mask"
1358
1359
1360
1361

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"
    def load_image(self, image, channel):
1362
        image_path = folder_paths.get_annotated_filepath(image)
1363
        i = Image.open(image_path)
1364
        i = ImageOps.exif_transpose(i)
1365
1366
        if i.getbands() != ("R", "G", "B", "A"):
            i = i.convert("RGBA")
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
        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):
1380
        image_path = folder_paths.get_annotated_filepath(image)
1381
1382
1383
1384
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
pythongosssss's avatar
pythongosssss committed
1385

1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
    @classmethod
    def VALIDATE_INPUTS(s, image, channel):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

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

        return True

comfyanonymous's avatar
comfyanonymous committed
1396
class ImageScale:
comfyanonymous's avatar
comfyanonymous committed
1397
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic"]
comfyanonymous's avatar
comfyanonymous committed
1398
1399
1400
1401
1402
    crop_methods = ["disabled", "center"]

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
1403
1404
                              "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
1405
1406
1407
1408
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

1409
    CATEGORY = "image/upscaling"
1410

comfyanonymous's avatar
comfyanonymous committed
1411
1412
    def upscale(self, image, upscale_method, width, height, crop):
        samples = image.movedim(-1,1)
comfyanonymous's avatar
comfyanonymous committed
1413
        s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
comfyanonymous's avatar
comfyanonymous committed
1414
1415
        s = s.movedim(1,-1)
        return (s,)
comfyanonymous's avatar
comfyanonymous committed
1416

comfyanonymous's avatar
comfyanonymous committed
1417
class ImageScaleBy:
comfyanonymous's avatar
comfyanonymous committed
1418
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic"]
comfyanonymous's avatar
comfyanonymous committed
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436

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

    CATEGORY = "image/upscaling"

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

1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
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,)

1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
class ImageBatch:

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

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

    CATEGORY = "image"

    def batch(self, image1, image2):
        if image1.shape[1:] != image2.shape[1:]:
            image2 = comfy.utils.common_upscale(image2.movedim(-1,1), image1.shape[2], image1.shape[1], "bilinear", "center").movedim(1,-1)
        s = torch.cat((image1, image2), dim=0)
        return (s,)
1468

comfyanonymous's avatar
comfyanonymous committed
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
class EmptyImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
                              "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
                              "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
                              }}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"

    CATEGORY = "image"

    def generate(self, width, height, batch_size=1, color=0):
        r = torch.full([batch_size, height, width, 1], ((color >> 16) & 0xFF) / 0xFF)
        g = torch.full([batch_size, height, width, 1], ((color >> 8) & 0xFF) / 0xFF)
        b = torch.full([batch_size, height, width, 1], ((color) & 0xFF) / 0xFF)
        return (torch.cat((r, g, b), dim=-1), )

Guo Y.K's avatar
Guo Y.K committed
1491
1492
1493
1494
1495
1496
1497
class ImagePadForOutpaint:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
1498
1499
1500
1501
                "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
                "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
1502
                "feathering": ("INT", {"default": 40, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
Guo Y.K's avatar
Guo Y.K committed
1503
1504
1505
1506
1507
1508
1509
1510
            }
        }

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

    CATEGORY = "image"

1511
    def expand_image(self, image, left, top, right, bottom, feathering):
Guo Y.K's avatar
Guo Y.K committed
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
        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,
        )
1524

1525
1526
1527
1528
1529
        t = torch.zeros(
            (d2, d3),
            dtype=torch.float32
        )

1530
        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549

            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
1550

Guo Y.K's avatar
Guo Y.K committed
1551
1552
1553
        return (new_image, mask)


comfyanonymous's avatar
comfyanonymous committed
1554
1555
NODE_CLASS_MAPPINGS = {
    "KSampler": KSampler,
1556
    "CheckpointLoaderSimple": CheckpointLoaderSimple,
comfyanonymous's avatar
comfyanonymous committed
1557
    "CLIPTextEncode": CLIPTextEncode,
comfyanonymous's avatar
comfyanonymous committed
1558
    "CLIPSetLastLayer": CLIPSetLastLayer,
comfyanonymous's avatar
comfyanonymous committed
1559
1560
    "VAEDecode": VAEDecode,
    "VAEEncode": VAEEncode,
1561
    "VAEEncodeForInpaint": VAEEncodeForInpaint,
comfyanonymous's avatar
comfyanonymous committed
1562
1563
1564
    "VAELoader": VAELoader,
    "EmptyLatentImage": EmptyLatentImage,
    "LatentUpscale": LatentUpscale,
comfyanonymous's avatar
comfyanonymous committed
1565
    "LatentUpscaleBy": LatentUpscaleBy,
1566
    "LatentFromBatch": LatentFromBatch,
1567
    "RepeatLatentBatch": RepeatLatentBatch,
comfyanonymous's avatar
comfyanonymous committed
1568
    "SaveImage": SaveImage,
pythongosssss's avatar
pythongosssss committed
1569
    "PreviewImage": PreviewImage,
comfyanonymous's avatar
comfyanonymous committed
1570
    "LoadImage": LoadImage,
1571
    "LoadImageMask": LoadImageMask,
comfyanonymous's avatar
comfyanonymous committed
1572
    "ImageScale": ImageScale,
comfyanonymous's avatar
comfyanonymous committed
1573
    "ImageScaleBy": ImageScaleBy,
1574
    "ImageInvert": ImageInvert,
1575
    "ImageBatch": ImageBatch,
Guo Y.K's avatar
Guo Y.K committed
1576
    "ImagePadForOutpaint": ImagePadForOutpaint,
comfyanonymous's avatar
comfyanonymous committed
1577
    "EmptyImage": EmptyImage,
FizzleDorf's avatar
FizzleDorf committed
1578
    "ConditioningAverage ": ConditioningAverage ,
comfyanonymous's avatar
comfyanonymous committed
1579
    "ConditioningCombine": ConditioningCombine,
1580
    "ConditioningConcat": ConditioningConcat,
comfyanonymous's avatar
comfyanonymous committed
1581
    "ConditioningSetArea": ConditioningSetArea,
Jacob Segal's avatar
Jacob Segal committed
1582
    "ConditioningSetMask": ConditioningSetMask,
comfyanonymous's avatar
comfyanonymous committed
1583
    "KSamplerAdvanced": KSamplerAdvanced,
1584
    "SetLatentNoiseMask": SetLatentNoiseMask,
comfyanonymous's avatar
comfyanonymous committed
1585
    "LatentComposite": LatentComposite,
1586
    "LatentBlend": LatentBlend,
comfyanonymous's avatar
comfyanonymous committed
1587
    "LatentRotate": LatentRotate,
comfyanonymous's avatar
comfyanonymous committed
1588
    "LatentFlip": LatentFlip,
comfyanonymous's avatar
comfyanonymous committed
1589
    "LatentCrop": LatentCrop,
1590
    "LoraLoader": LoraLoader,
1591
    "CLIPLoader": CLIPLoader,
1592
    "UNETLoader": UNETLoader,
1593
    "DualCLIPLoader": DualCLIPLoader,
1594
    "CLIPVisionEncode": CLIPVisionEncode,
1595
    "StyleModelApply": StyleModelApply,
1596
    "unCLIPConditioning": unCLIPConditioning,
comfyanonymous's avatar
comfyanonymous committed
1597
    "ControlNetApply": ControlNetApply,
1598
    "ControlNetApplyAdvanced": ControlNetApplyAdvanced,
comfyanonymous's avatar
comfyanonymous committed
1599
    "ControlNetLoader": ControlNetLoader,
1600
    "DiffControlNetLoader": DiffControlNetLoader,
comfyanonymous's avatar
comfyanonymous committed
1601
1602
    "StyleModelLoader": StyleModelLoader,
    "CLIPVisionLoader": CLIPVisionLoader,
1603
    "VAEDecodeTiled": VAEDecodeTiled,
comfyanonymous's avatar
comfyanonymous committed
1604
    "VAEEncodeTiled": VAEEncodeTiled,
1605
    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
1606
1607
1608
    "GLIGENLoader": GLIGENLoader,
    "GLIGENTextBoxApply": GLIGENTextBoxApply,

1609
    "CheckpointLoader": CheckpointLoader,
sALTaccount's avatar
sALTaccount committed
1610
    "DiffusersLoader": DiffusersLoader,
Dr.Lt.Data's avatar
Dr.Lt.Data committed
1611
1612

    "LoadLatent": LoadLatent,
1613
    "SaveLatent": SaveLatent,
1614
1615

    "ConditioningZeroOut": ConditioningZeroOut,
1616
    "ConditioningSetTimestepRange": ConditioningSetTimestepRange,
comfyanonymous's avatar
comfyanonymous committed
1617
1618
}

City's avatar
City committed
1619
1620
1621
1622
1623
NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "KSampler": "KSampler",
    "KSamplerAdvanced": "KSampler (Advanced)",
    # Loaders
comfyanonymous's avatar
Rename.  
comfyanonymous committed
1624
1625
    "CheckpointLoader": "Load Checkpoint (With Config)",
    "CheckpointLoaderSimple": "Load Checkpoint",
City's avatar
City committed
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
    "VAELoader": "Load VAE",
    "LoraLoader": "Load LoRA",
    "CLIPLoader": "Load CLIP",
    "ControlNetLoader": "Load ControlNet Model",
    "DiffControlNetLoader": "Load ControlNet Model (diff)",
    "StyleModelLoader": "Load Style Model",
    "CLIPVisionLoader": "Load CLIP Vision",
    "UpscaleModelLoader": "Load Upscale Model",
    # Conditioning
    "CLIPVisionEncode": "CLIP Vision Encode",
    "StyleModelApply": "Apply Style Model",
    "CLIPTextEncode": "CLIP Text Encode (Prompt)",
    "CLIPSetLastLayer": "CLIP Set Last Layer",
    "ConditioningCombine": "Conditioning (Combine)",
FizzleDorf's avatar
FizzleDorf committed
1640
    "ConditioningAverage ": "Conditioning (Average)",
1641
    "ConditioningConcat": "Conditioning (Concat)",
City's avatar
City committed
1642
    "ConditioningSetArea": "Conditioning (Set Area)",
Jacob Segal's avatar
Jacob Segal committed
1643
    "ConditioningSetMask": "Conditioning (Set Mask)",
City's avatar
City committed
1644
    "ControlNetApply": "Apply ControlNet",
1645
    "ControlNetApplyAdvanced": "Apply ControlNet (Advanced)",
City's avatar
City committed
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
    # Latent
    "VAEEncodeForInpaint": "VAE Encode (for Inpainting)",
    "SetLatentNoiseMask": "Set Latent Noise Mask",
    "VAEDecode": "VAE Decode",
    "VAEEncode": "VAE Encode",
    "LatentRotate": "Rotate Latent",
    "LatentFlip": "Flip Latent",
    "LatentCrop": "Crop Latent",
    "EmptyLatentImage": "Empty Latent Image",
    "LatentUpscale": "Upscale Latent",
comfyanonymous's avatar
comfyanonymous committed
1656
    "LatentUpscaleBy": "Upscale Latent By",
City's avatar
City committed
1657
    "LatentComposite": "Latent Composite",
1658
    "LatentBlend": "Latent Blend",
1659
1660
    "LatentFromBatch" : "Latent From Batch",
    "RepeatLatentBatch": "Repeat Latent Batch",
City's avatar
City committed
1661
1662
1663
1664
1665
1666
    # Image
    "SaveImage": "Save Image",
    "PreviewImage": "Preview Image",
    "LoadImage": "Load Image",
    "LoadImageMask": "Load Image (as Mask)",
    "ImageScale": "Upscale Image",
comfyanonymous's avatar
comfyanonymous committed
1667
    "ImageScaleBy": "Upscale Image By",
City's avatar
City committed
1668
1669
1670
    "ImageUpscaleWithModel": "Upscale Image (using Model)",
    "ImageInvert": "Invert Image",
    "ImagePadForOutpaint": "Pad Image for Outpainting",
1671
    "ImageBatch": "Batch Images",
City's avatar
City committed
1672
1673
1674
1675
1676
    # _for_testing
    "VAEDecodeTiled": "VAE Decode (Tiled)",
    "VAEEncodeTiled": "VAE Encode (Tiled)",
}

1677
1678
EXTENSION_WEB_DIRS = {}

1679
def load_custom_node(module_path, ignore=set()):
1680
1681
1682
1683
1684
1685
1686
    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)
1687
            module_dir = os.path.split(module_path)[0]
1688
1689
        else:
            module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
1690
1691
            module_dir = module_path

1692
1693
1694
        module = importlib.util.module_from_spec(module_spec)
        sys.modules[module_name] = module
        module_spec.loader.exec_module(module)
1695
1696
1697
1698
1699
1700

        if hasattr(module, "WEB_DIRECTORY") and getattr(module, "WEB_DIRECTORY") is not None:
            web_dir = os.path.abspath(os.path.join(module_dir, getattr(module, "WEB_DIRECTORY")))
            if os.path.isdir(web_dir):
                EXTENSION_WEB_DIRS[module_name] = web_dir

1701
        if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
1702
1703
1704
            for name in module.NODE_CLASS_MAPPINGS:
                if name not in ignore:
                    NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name]
1705
1706
            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)
1707
            return True
1708
1709
        else:
            print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
1710
            return False
1711
1712
1713
    except Exception as e:
        print(traceback.format_exc())
        print(f"Cannot import {module_path} module for custom nodes:", e)
1714
        return False
1715

Hacker 17082006's avatar
Hacker 17082006 committed
1716
def load_custom_nodes():
1717
    base_node_names = set(NODE_CLASS_MAPPINGS.keys())
1718
    node_paths = folder_paths.get_folder_paths("custom_nodes")
1719
    node_import_times = []
1720
1721
1722
1723
1724
1725
1726
1727
    for custom_node_path in node_paths:
        possible_modules = os.listdir(custom_node_path)
        if "__pycache__" in possible_modules:
            possible_modules.remove("__pycache__")

        for possible_module in possible_modules:
            module_path = os.path.join(custom_node_path, possible_module)
            if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
1728
            if module_path.endswith(".disabled"): continue
1729
            time_before = time.perf_counter()
1730
            success = load_custom_node(module_path, base_node_names)
1731
            node_import_times.append((time.perf_counter() - time_before, module_path, success))
1732

1733
    if len(node_import_times) > 0:
comfyanonymous's avatar
comfyanonymous committed
1734
        print("\nImport times for custom nodes:")
1735
        for n in sorted(node_import_times):
1736
1737
1738
1739
1740
            if n[2]:
                import_message = ""
            else:
                import_message = " (IMPORT FAILED)"
            print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
1741
        print()
1742

1743
def init_custom_nodes():
1744
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_hypernetwork.py"))
1745
1746
    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"))
1747
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))
1748
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_rebatch.py"))
1749
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_model_merging.py"))
1750
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_tomesd.py"))
1751
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_clip_sdxl.py"))
1752
    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_canny.py"))
1753
    load_custom_nodes()