From 2a87f6630d5181bea101c5986ee1d64d59c11c36 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 26 Jan 2023 12:23:15 -0500
Subject: [PATCH] Some categories for the nodes.

---
 nodes.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/nodes.py b/nodes.py
index 11a78112..cb69e1a8 100644
--- a/nodes.py
+++ b/nodes.py
@@ -33,6 +33,8 @@ class CLIPTextEncode:
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "encode"
 
+    CATEGORY = "conditioning"
+
     def encode(self, clip, text):
         return ([[clip.encode(text), {}]], )
 
@@ -43,6 +45,8 @@ class ConditioningCombine:
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "combine"
 
+    CATEGORY = "conditioning"
+
     def combine(self, conditioning_1, conditioning_2):
         return (conditioning_1 + conditioning_2, )
 
@@ -59,6 +63,8 @@ class ConditioningSetArea:
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "append"
 
+    CATEGORY = "conditioning"
+
     def append(self, conditioning, width, height, x, y, strength, min_sigma=0.0, max_sigma=99.0):
         c = copy.deepcopy(conditioning)
         for t in c:
@@ -78,6 +84,8 @@ class VAEDecode:
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "decode"
 
+    CATEGORY = "latent"
+
     def decode(self, vae, samples):
         return (vae.decode(samples), )
 
@@ -91,6 +99,8 @@ class VAEEncode:
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "encode"
 
+    CATEGORY = "latent"
+
     def encode(self, vae, pixels):
         x = (pixels.shape[1] // 64) * 64
         y = (pixels.shape[2] // 64) * 64
@@ -110,6 +120,8 @@ class CheckpointLoader:
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
     FUNCTION = "load_checkpoint"
 
+    CATEGORY = "loaders"
+
     def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
         config_path = os.path.join(self.config_dir, config_name)
         ckpt_path = os.path.join(self.ckpt_dir, ckpt_name)
@@ -124,6 +136,8 @@ class VAELoader:
     RETURN_TYPES = ("VAE",)
     FUNCTION = "load_vae"
 
+    CATEGORY = "loaders"
+
     #TODO: scale factor?
     def load_vae(self, vae_name):
         vae_path = os.path.join(self.vae_dir, vae_name)
@@ -142,6 +156,8 @@ class EmptyLatentImage:
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "generate"
 
+    CATEGORY = "latent"
+
     def generate(self, width, height, batch_size=1):
         latent = torch.zeros([batch_size, 4, height // 8, width // 8])
         return (latent, )
@@ -199,6 +215,8 @@ class KSampler:
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "sample"
 
+    CATEGORY = "sampling"
+
     def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
         noise = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=torch.manual_seed(seed), device="cpu")
         model = model.to(self.device)
@@ -250,6 +268,8 @@ class SaveImage:
 
     OUTPUT_NODE = True
 
+    CATEGORY = "image"
+
     def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         def map_filename(filename):
             prefix_len = len(filename_prefix)
@@ -283,6 +303,8 @@ class LoadImage:
                     {"image": (os.listdir(s.input_dir), )},
                 }
 
+    CATEGORY = "image"
+
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "load_image"
     def load_image(self, image):