From 0c2963062a064c20e0b5aed6900829305fb370f9 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 3 Dec 2023 16:11:49 +0200
Subject: [PATCH] Add label node

---
 nodes.py | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index b920f76..3d4b415 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2554,6 +2554,58 @@ class InjectNoiseToLatent:
         
         samples["samples"] = noised
         return (samples,)
+
+class AddLabel:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "image":("IMAGE",),  
+            "text_x": ("INT", {"default": 10, "min": 0, "max": 4096, "step": 1}),
+            "text_y": ("INT", {"default": 2, "min": 0, "max": 4096, "step": 1}),
+            "height": ("INT", {"default": 48, "min": 0, "max": 4096, "step": 1}),
+            "font_size": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 1}),
+            "font_color": ("STRING", {"default": "white"}),
+            "label_color": ("STRING", {"default": "black"}),
+            "font_path": ("STRING", {"default": "fonts\\TTNorms-Black.otf"}),
+            "text": ("STRING", {"default": "Text"}),
+            "direction": (
+            [   'up',
+                'down',
+            ],
+            {
+            "default": 'up'
+            
+             }),
+            },
+            }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "addlabel"
+
+    CATEGORY = "KJNodes"
+        
+    def addlabel(self, image, text_x, text_y, text, height, font_size, font_color, label_color, font_path, direction):
+        batch_size = image.shape[0]
+        width = image.shape[2]
+        
+        if font_path == "fonts\\TTNorms-Black.otf": #I don't know why relative path won't work otherwise...
+            font_path = os.path.join(script_dir, font_path)
+            
+        label_image = Image.new("RGB", (width, height), label_color)
+        draw = ImageDraw.Draw(label_image)
+        font = ImageFont.truetype(font_path, font_size)
+        draw.text((text_x, text_y), text, font=font, fill=font_color) 
+        label_image = np.array(label_image).astype(np.float32) / 255.0
+        label_image = torch.from_numpy(label_image)[None, :, :, :]
+        # Duplicate the label image for the entire batch
+        label_batch = label_image.repeat(batch_size, 1, 1, 1)
+
+        if direction == 'down':
+            combined_images = torch.cat((image, label_batch), dim=1)
+        elif direction == 'up':
+            combined_images = torch.cat((label_batch, image), dim=1)
+        
+        return (combined_images,)
     
 NODE_CLASS_MAPPINGS = {
     "INTConstant": INTConstant,
@@ -2602,7 +2654,8 @@ NODE_CLASS_MAPPINGS = {
     "DummyLatentOut": DummyLatentOut,
     "NormalizeLatent": NormalizeLatent,
     "FlipSigmasAdjusted": FlipSigmasAdjusted,
-    "InjectNoiseToLatent": InjectNoiseToLatent
+    "InjectNoiseToLatent": InjectNoiseToLatent,
+    "AddLabel": AddLabel
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "INTConstant": "INT Constant",
@@ -2650,6 +2703,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DummyLatentOut": "DummyLatentOut",
     "NormalizeLatent": "NormalizeLatent",
     "FlipSigmasAdjusted": "FlipSigmasAdjusted",
-    "InjectNoiseToLatent": "InjectNoiseToLatent"
+    "InjectNoiseToLatent": "InjectNoiseToLatent",
+    "AddLabel": "AddLabel"
 
 }
\ No newline at end of file