adding the overlay node, adjust node names

MohammadAboulEla · Aug 16, 2024 · 2c43f78 · 2c43f78
1 parent 4c0c481
commit 2c43f78
Show file tree

Hide file tree

Showing 7 changed files with 167 additions and 82 deletions.
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 The iTools are some quality of life nodes, like read a possible prompt used to create an image, save a prompt to file as a new line, read prompts from a multiline file.
 
 ### Nodes:
-**iTools ImageLoaderPlus** is an enhancement of the original ComfyUI ImageLoader node. It attempts to return the possible prompt used to create an image. If the prompt isn't found directly, this node will search for the prompt within the following supported nodes in the workflow:
+**iTools Image Loader Plus** is an enhancement of the original ComfyUI ImageLoader node. It attempts to return the possible prompt used to create an image. If the prompt isn't found directly, this node will search for the prompt within the following supported nodes in the workflow:
 
   - `CLIPTextEncodeSDXL`
   - `CLIPTextEncode`
@@ -13,10 +13,16 @@ The iTools are some quality of life nodes, like read a possible prompt used to c
 
 If your prompt is within any of these nodes, you will be able to retrieve it. The `ShowText|pysssss` node is a convenient addition that you can integrate it into any workflow to make sure it will be retrieved by `iTools ImageLoaderPlus`.
 
-![iTools ImageLoaderPlus](https://github.com/MohammadAboulEla/ComfyUI-iTools/blob/master/examples/Screenshot1.png)
+![iTools ImageLoaderPlus](https://github.com/MohammadAboulEla/ComfyUI-iTools/blob/master/examples/Screenshot1.jpg)
 
-**iTools PromptLoader:** will try to load a prompt from a txt file, you need to provide a full path to a .txt file or try use the default prompt.txt one is the examples folder, the node will return the prompt at given index (line number) not that count start from zero, node also will return a radom prompt from that file.
+**iTools Prompt Loader:** will try to load a prompt from a txt file, you need to provide a full path to a .txt file or try use the default prompt.txt one is the examples folder, the node will return the prompt at given index (line number) not that count start from zero, node also will return a radom prompt from that file.
 
-**iTools PromptSaver:** will append the given prompt as a new line to the given file, provide a full path to a .txt file or try use the default prompt.txt one.
+**iTools Prompt Saver:** will append the given prompt as a new line to the given file, provide a full path to a .txt file or try use the default prompt.txt one.
 
 [![Watch the video](https://github.com/user-attachments/assets/22af7830-066f-498e-a90f-0513b56fa343)](https://github.com/user-attachments/assets/22af7830-066f-498e-a90f-0513b56fa343)
+
+**iTools Add Text Overlay:** will add an overlay bottom bar to show a given text, you may change the background color of the overlay bar and the fon size
+
+![iTools Add Text Overlay](https://github.com/MohammadAboulEla/ComfyUI-iTools/blob/master/examples/Screenshot2.jpg)
+
+![iTools Add Text Overlay](https://github.com/MohammadAboulEla/ComfyUI-iTools/blob/master/examples/Screenshot3.jpg)
diff --git a/examples/Screenshot1.png b/examples/Screenshot1.png
diff --git a/examples/prompts.txt b/examples/prompts.txt
@@ -5,3 +5,4 @@ Five kids playing with sticks
 An adorable ladybug flamingo in a field of pink and red flowers
 Pears and peaches some open and some falling, with a splash of ice water and ice in show freshness of both fruitscute cat in the forest
 Cute cat in the forest
+A glass jar that is filled with tiny vehicles (cars, trucks, motorcycles). Fanciful sticker label on the jar says: "Traffic Jam". Advertisement photography shot
diff --git a/iTools_nodes.py b/iTools_nodes.py
@@ -5,83 +5,14 @@
 import torch
 import numpy as np
 import hashlib
+import torchvision.transforms.v2 as T
 
 import node_helpers
 from PIL import Image, ImageSequence, ImageOps
 from .metadata.metadata_extractor import get_prompt
 from .metadata.file_handeler import FileHandler
+from .metadata.overlay import add_overlay_bar
 
-class IToolsPromptLoader:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {
-                        "file_path": ("STRING", {"default": 'prompts.txt',"multiline": False}),
-                        "seed": ("INT", {"default": 0,"control_after_generate":0, "min": 0, "max": 0xfffff})
-                     }
-                }
-
-    CATEGORY = "iTools"
-
-    RETURN_TYPES = ("STRING","STRING",)
-    RETURN_NAMES = ("Prompt at index","Prompt random")
-    FUNCTION = "load_file"
-
-    def load_file(self,file_path, seed):
-        prompt= ""
-        prompt_random = ""
-        cn = folder_paths.folder_names_and_paths["custom_nodes"][0][0]
-        if file_path == "prompts.txt":
-            file = os.path.join(cn, "ComfyUi-iTools\examples\prompts.txt")
-        else:
-            file = file_path.replace('"','')
-        if os.path.exists(file):
-            fh = FileHandler(file)
-            try:
-                line = fh.read_line(seed)
-                prompt = fh.unescape_quotes(line)
-            except IndexError:
-                prompt = f"IndexError: Line {seed} is not available"
-            try:
-                r_index = random.randint(0,fh.len_lines()-1)
-                prompt_random = fh.unescape_quotes(fh.read_line(r_index))
-            except IndexError:
-                prompt = f"IndexError: Line {r_index} is not available"
-        else:
-            prompt = f"File not exist, {file}"
-        return (prompt, prompt_random)
-
-class IToolsPromptSaver:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {   "prompt": ("STRING", {"forceInput": True}),
-                        "file_path": ("STRING", {"default": 'prompts.txt',"multiline": False}),
-                     }
-                }
-
-    CATEGORY = "iTools"
-    RETURN_TYPES = ()
-    OUTPUT_NODE = True
-    FUNCTION = "save_to_file"
-
-    def save_to_file(self,file_path, prompt):
-        cn = folder_paths.folder_names_and_paths["custom_nodes"][0][0]
-        if file_path == "prompts.txt":
-            file = os.path.join(cn, "ComfyUi-iTools\examples\prompts.txt")
-        else:
-            file = file_path.replace('"','')
-        if os.path.exists(file) and prompt is not None and prompt != "":
-            fh = FileHandler(file)
-            try:
-                fh.append_line(prompt)
-                print(f"Prompt: {prompt} saved to {file}")
-            except Exception as e:
-                print(f"Error while writing the prompt: {e}")
-        else:
-            print(f"Error while writing the prompt: {e}")
-        return (True,)
 
 class IToolsLoadImagePlus:
     @classmethod
@@ -159,18 +90,137 @@ def VALIDATE_INPUTS(s, image):
         return True
 
 
+class IToolsPromptLoader:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+            {
+                "file_path": ("STRING", {"default": 'prompts.txt', "multiline": False}),
+                "seed": ("INT", {"default": 0, "control_after_generate": 0, "min": 0, "max": 0xfffff})
+            }
+        }
+
+    CATEGORY = "iTools"
+
+    RETURN_TYPES = ("STRING", "STRING",)
+    RETURN_NAMES = ("Prompt at index", "Prompt random")
+    FUNCTION = "load_file"
+
+    def load_file(self, file_path, seed):
+        prompt = ""
+        prompt_random = ""
+        cn = folder_paths.folder_names_and_paths["custom_nodes"][0][0]
+        if file_path == "prompts.txt":
+            file = os.path.join(cn, "ComfyUi-iTools\examples\prompts.txt")
+        else:
+            file = file_path.replace('"', '')
+        if os.path.exists(file):
+            fh = FileHandler(file)
+            try:
+                line = fh.read_line(seed)
+                prompt = fh.unescape_quotes(line)
+            except IndexError:
+                prompt = f"IndexError: Line {seed} is not available"
+            try:
+                r_index = random.randint(0, fh.len_lines() - 1)
+                prompt_random = fh.unescape_quotes(fh.read_line(r_index))
+            except IndexError:
+                prompt = f"IndexError: Line {r_index} is not available"
+        else:
+            prompt = f"File not exist, {file}"
+        return (prompt, prompt_random)
+
+
+class IToolsPromptSaver:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"prompt": ("STRING", {"forceInput": True}),
+                     "file_path": ("STRING", {"default": 'prompts.txt', "multiline": False}),
+                     }
+                }
+
+    CATEGORY = "iTools"
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "save_to_file"
+
+    def save_to_file(self, file_path, prompt):
+        cn = folder_paths.folder_names_and_paths["custom_nodes"][0][0]
+        if file_path == "prompts.txt":
+            file = os.path.join(cn, "ComfyUi-iTools\examples\prompts.txt")
+        else:
+            file = file_path.replace('"', '')
+        if os.path.exists(file) and prompt is not None and prompt != "":
+            fh = FileHandler(file)
+            try:
+                fh.append_line(prompt)
+                print(f"Prompt: {prompt} saved to {file}")
+            except Exception as e:
+                print(f"Error while writing the prompt: {e}")
+        else:
+            print(f"Error while writing the prompt")
+        return (True,)
+
+
+class IToolsAddOverlay:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required":
+            {
+                "image": ("IMAGE", {}),
+                "text": ("STRING", {"default": 'img info:', "multiline": False}),
+                "background_color": ("STRING", {"default": '#000000AA', "multiline": False}),
+                "font_size": ("INT", {"default": 40, "min": 10, "max": 1000})
+            }
+        }
+
+    CATEGORY = "iTools"
+    RETURN_TYPES = ("IMAGE",)
+    # OUTPUT_NODE = True
+    FUNCTION = "add_text_overlay"
+
+    def add_text_overlay(self, image, text, font_size, background_color):
+        # Remove the batch dimension and rearrange to [C, H, W]
+        tensor = image.squeeze(0).permute(2, 0, 1)
+
+        # Ensure the values are in the range [0, 1]
+        tensor = tensor.clamp(0, 1)
+
+        # Convert to PIL Image
+        to_pil = T.ToPILImage()
+        pil_image = to_pil(tensor)
+
+        # Add overlay (assuming add_overlay_bar is defined elsewhere)
+        composite = add_overlay_bar(pil_image, text, font_size=font_size, background_color=background_color)
+
+        # Convert back to tensor
+        to_tensor = T.ToTensor()
+        out = to_tensor(composite)
+
+        # Add batch dimension to match original input
+        out = out.unsqueeze(0)
+
+        # Rearrange back to [B, H, W, C] to match input format
+        out = out.permute(0, 2, 3, 1)
+
+        return (out,)
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "iToolsLoadImagePlus": IToolsLoadImagePlus,
-    "iToolsPromptLoader" : IToolsPromptLoader,
-    "iToolsPromptSaver" : IToolsPromptSaver
+    "iToolsPromptLoader": IToolsPromptLoader,
+    "iToolsPromptSaver": IToolsPromptSaver,
+    "iToolsAddOverlay": IToolsAddOverlay
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "iToolsLoadImagePlus": "iTools LoadImagePlus",
-    "iToolsPromptLoader" : "iTools PromptLoader",
-    "iToolsPromptSaver" : "iTools PromptSaver"
+    "iToolsLoadImagePlus": "iTools Load Image Plus",
+    "iToolsPromptLoader": "iTools Prompt Loader",
+    "iToolsPromptSaver": "iTools Prompt Saver",
+    "iToolsAddOverlay": "iTools Add Text Overlay"
 }
-
diff --git a/metadata/Inconsolata.otf b/metadata/Inconsolata.otf
diff --git a/metadata/overlay.py b/metadata/overlay.py
@@ -0,0 +1,28 @@
+
+import os
+from PIL import Image, ImageDraw, ImageFont, ImageColor
+
+
+def add_overlay_bar(image, info, font_size=32, background_color="#000000AA"):
+    font_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "Inconsolata.otf")
+    font = ImageFont.truetype(font_path, font_size)
+    txt_height = font.getmask("Q").getbbox()[3] + font.getmetrics()[1] + 5
+    try:
+        background_color = ImageColor.getrgb(background_color)
+    except ValueError:
+        print(f"[iTools] Unknown color specifier: {background_color} will use default one")
+        background_color = ImageColor.getrgb("#000000AA")
+
+    txt = Image.new('RGBA', (image.width, txt_height), color=background_color)
+    draw = ImageDraw.Draw(txt)
+    draw.text((0, 0), f"{info}", font=font, fill=(255, 255, 255, 255))
+
+    composite = Image.new('RGBA', image.size)
+    composite.paste(image, (0, 0))
+    composite.paste(txt, (0, image.height - txt_height), mask=txt)
+
+    return composite
+
+
+if __name__ == '__main__':
+    ...
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-itools"
 description = "The iTools are some quality of life nodes, like read a possible prompt used to create an image, save a prompt to file as a new line, read prompts from a multiline file."
-version = "0.2.2"
+version = "0.3.2"
 license = {file = "LICENSE"}
 
 [project.urls]