From 40d4853393652bdd12f0693145b18bf64aa370e6 Mon Sep 17 00:00:00 2001 From: Fillip Date: Sun, 1 Sep 2024 22:45:32 -0700 Subject: [PATCH] added ollama captioning tool --- __init__.py | 12 +++- nodes/FL_ColorPicker.py | 2 +- nodes/FL_GradGen.py | 4 +- nodes/FL_OllamaCaptioner.py | 100 ++++++++++++++++++++++++++++ nodes/SaveAndDisplayImage.py | 29 ++++++++ web/nodes/GradientImageGenerator.js | 4 +- web/nodes/SaveAndDisplayImage.js | 10 +++ 7 files changed, 153 insertions(+), 8 deletions(-) create mode 100644 nodes/FL_OllamaCaptioner.py create mode 100644 nodes/SaveAndDisplayImage.py create mode 100644 web/nodes/SaveAndDisplayImage.js diff --git a/__init__.py b/__init__.py index c51d5da..f8f5f03 100644 --- a/__init__.py +++ b/__init__.py @@ -34,7 +34,7 @@ from .nodes.FL_TetrisGame import FL_TetrisGame from .nodes.FL_Dither import FL_Dither from .nodes.FL_SystemCheck import FL_SystemCheck from .nodes.FL_ColorPicker import FL_ColorPicker -from .nodes.FL_GradGen import GradientImageGenerator +from .nodes.FL_GradGen import FL_GradGenerator from .nodes.FL_MirrorAndAppendCaptions import FL_MirrorAndAppendCaptions from .nodes.FL_ImageCaptionLayout import FL_ImageCaptionLayout from .nodes.FL_HFHubModelUploader import FL_HFHubModelUploader @@ -66,6 +66,8 @@ from .nodes.FL_PDFMerger import FL_PDFMerger from .nodes.FL_PDFTextExtractor import FL_PDFTextExtractor from .nodes.FL_PDFImageExtractor import FL_PDFImageExtractor from .nodes.FL_BulkPDFLoader import FL_BulkPDFLoader +from .nodes.SaveAndDisplayImage import SaveAndDisplayImage +from .nodes.FL_OllamaCaptioner import FL_OllamaCaptioner @@ -107,7 +109,7 @@ NODE_CLASS_MAPPINGS = { "FL_Dither": FL_Dither, "FL_SystemCheck": FL_SystemCheck, "FL_ColorPicker": FL_ColorPicker, - "GradientImageGenerator": GradientImageGenerator, + "FL_GradGenerator": FL_GradGenerator, "FL_MirrorAndAppendCaptions": FL_MirrorAndAppendCaptions, "FL_ImageCaptionLayout": FL_ImageCaptionLayout, "FL_HFHubModelUploader": FL_HFHubModelUploader, @@ -139,6 +141,8 @@ NODE_CLASS_MAPPINGS = { "FL_PDFTextExtractor": FL_PDFTextExtractor, "FL_PDFImageExtractor": FL_PDFImageExtractor, "FL_BulkPDFLoader": FL_BulkPDFLoader, + "SaveAndDisplayImage": SaveAndDisplayImage, + "FL_OllamaCaptioner": FL_OllamaCaptioner, } @@ -179,7 +183,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FL_Dither": "FL Dither", "FL_SystemCheck": "FL System Check", "FL_ColorPicker": "FL Color Picker", - "GradientImageGenerator": "GradientImageGenerator", + "FL_GradGenerator": "FL Grad Generator", "FL_MirrorAndAppendCaptions": "FL Mirror And Append Captions", "FL_ImageCaptionLayout": "FL Image Caption Layout", "FL_HFHubModelUploader": "FL HFHub Model Uploader", @@ -211,6 +215,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FL_PDFTextExtractor": "FL PDF Text Extractor", "FL_PDFImageExtractor": "FL PDF Image Extractor", "FL_BulkPDFLoader": "FL Bulk PDF Loader", + "SaveAndDisplayImage": "Save And Display Image", + "FL_OllamaCaptioner": "FL Ollama Captioner by Cosmic", } diff --git a/nodes/FL_ColorPicker.py b/nodes/FL_ColorPicker.py index a0b0962..9af9182 100644 --- a/nodes/FL_ColorPicker.py +++ b/nodes/FL_ColorPicker.py @@ -9,7 +9,7 @@ class FL_ColorPicker: RETURN_TYPES = ("STRING",) FUNCTION = "get_color" - CATEGORY = "ui" + CATEGORY = "🏵️Fill Nodes/experiments" def get_color(self, selected_color): return (selected_color,) \ No newline at end of file diff --git a/nodes/FL_GradGen.py b/nodes/FL_GradGen.py index f5c3380..04660fb 100644 --- a/nodes/FL_GradGen.py +++ b/nodes/FL_GradGen.py @@ -4,7 +4,7 @@ import torch from PIL import Image -class GradientImageGenerator: +class FL_GradGenerator: @classmethod def INPUT_TYPES(s): return { @@ -76,7 +76,7 @@ class GradientImageGenerator: NODE_CLASS_MAPPINGS = { - "GradientImageGenerator": GradientImageGenerator + "GradientImageGenerator": FL_GradGenerator } NODE_DISPLAY_NAME_MAPPINGS = { diff --git a/nodes/FL_OllamaCaptioner.py b/nodes/FL_OllamaCaptioner.py new file mode 100644 index 0000000..56de557 --- /dev/null +++ b/nodes/FL_OllamaCaptioner.py @@ -0,0 +1,100 @@ +import os +import re +from PIL import Image +import numpy as np +from comfy.utils import ProgressBar +from ollama import Client +from io import BytesIO +import base64 + +class FL_OllamaCaptioner: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", {}), + "folder_name": ("STRING", {"default": "output_folder"}), + "use_llm": ("BOOLEAN", {"default": True}), + "url": ("STRING", {"default": "http://127.0.0.1:11434"}), # Default Ollama URL + "model": ("STRING", {"default": "default_model"}), # Replace with your model name + "overwrite": ("BOOLEAN", {"default": True}) + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "save_images_with_captions" + CATEGORY = "🏵️Fill Nodes/Captioning" + OUTPUT_NODE = True + + def sanitize_text(self, text): + return re.sub(r'[^a-zA-Z0-9\s.,!?-]', '', text) + + def generate_caption_with_ollama(self, image_tensor, url, model): + # Convert tensor to numpy array + image_np = image_tensor.cpu().numpy() + # Ensure the image is in the correct shape (height, width, channels) + if image_np.shape[0] == 1: # If the first dimension is 1, squeeze it + image_np = np.squeeze(image_np, axis=0) + if len(image_np.shape) == 2: + image_np = np.stack((image_np,) * 3, axis=-1) + elif image_np.shape[2] == 1: # If it's (height, width, 1) + image_np = np.repeat(image_np, 3, axis=2) + # Ensure values are in 0-255 range + image_np = (image_np * 255).clip(0, 255).astype(np.uint8) + # Convert to PIL Image + image = Image.fromarray(image_np) + + # Encode image to base64 + buffered = BytesIO() + image.save(buffered, format="PNG") + img_bytes = base64.b64encode(buffered.getvalue()).decode('utf-8') + + client = Client(host=url) + response = client.generate(model=model, prompt="describe the image", images=[img_bytes]) + + # Extract the caption from the response + return response['response'] + + def save_images_with_captions(self, images, folder_name, use_llm, url, model, overwrite): + os.makedirs(folder_name, exist_ok=True) + + saved_files = [] + pbar = ProgressBar(len(images)) + for i, image_tensor in enumerate(images): + base_name = f"image_{i}" + image_file_name = f"{folder_name}/{base_name}.png" + text_file_name = f"{folder_name}/{base_name}.txt" + + if not overwrite: + counter = 1 + while os.path.exists(image_file_name) or os.path.exists(text_file_name): + image_file_name = f"{folder_name}/{base_name}_{counter}.png" + text_file_name = f"{folder_name}/{base_name}_{counter}.txt" + counter += 1 + + if use_llm: + caption = self.generate_caption_with_ollama(image_tensor, url, model) + else: + caption = "Default Caption" + + sanitized_caption = self.sanitize_text(caption) + + # Convert tensor to numpy array and save the image as in the previous code + image_np = image_tensor.cpu().numpy() + if image_np.shape[0] == 1: + image_np = np.squeeze(image_np, axis=0) + if len(image_np.shape) == 2: + image_np = np.stack((image_np,) * 3, axis=-1) + elif image_np.shape[2] == 1: + image_np = np.repeat(image_np, 3, axis=2) + image_np = (image_np * 255).clip(0, 255).astype(np.uint8) + image = Image.fromarray(image_np) + image.save(image_file_name) + saved_files.append(image_file_name) + + with open(text_file_name, "w") as text_file: + text_file.write(sanitized_caption) + + pbar.update_absolute(i) + + return (f"Saved {len(images)} images and generated captions in '{folder_name}'",) diff --git a/nodes/SaveAndDisplayImage.py b/nodes/SaveAndDisplayImage.py new file mode 100644 index 0000000..0094c9c --- /dev/null +++ b/nodes/SaveAndDisplayImage.py @@ -0,0 +1,29 @@ +import os +import folder_paths +from nodes import SaveImage + + +class SaveAndDisplayImage(SaveImage): + @classmethod + def INPUT_TYPES(s): + return {"required": + {"images": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, + "hidden": + {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"} + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "save_and_display_images" + OUTPUT_NODE = True + CATEGORY = "image" + + def save_and_display_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + results = self.save_images(images, filename_prefix, prompt, extra_pnginfo) + + return { + "ui": { + "images": results['ui']['images'] + }, + "result": (images,) + } \ No newline at end of file diff --git a/web/nodes/GradientImageGenerator.js b/web/nodes/GradientImageGenerator.js index bb58d26..a7a4e21 100644 --- a/web/nodes/GradientImageGenerator.js +++ b/web/nodes/GradientImageGenerator.js @@ -1,9 +1,9 @@ import { app } from "../../../scripts/app.js"; app.registerExtension({ - name: "Comfy.GradientImageGenerator", + name: "Comfy.FL_GradGenerator", async nodeCreated(node) { - if (node.comfyClass === "GradientImageGenerator") { + if (node.comfyClass === "FL_GradGenerator") { addGradientImageGenerator(node); } } diff --git a/web/nodes/SaveAndDisplayImage.js b/web/nodes/SaveAndDisplayImage.js new file mode 100644 index 0000000..5e92258 --- /dev/null +++ b/web/nodes/SaveAndDisplayImage.js @@ -0,0 +1,10 @@ +import { app } from "../../../scripts/app.js"; + +app.registerExtension({ + name: "Comfy.SaveAndDisplayImage", + async nodeCreated(node) { + if (node.comfyClass === "SaveAndDisplayImage") { + node.addWidget("image", "preview", ""); + } + } +}); \ No newline at end of file