nodes/filliptm-ComfyUI_Fill-Nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added ollama captioning tool

Browse Source
This commit is contained in:
Fillip
2024-09-01 22:45:32 -07:00
parent 626b454918
commit 40d4853393
7 changed files with 153 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+9 -3
View File
@@ -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",
}
nodes/FL_ColorPicker.py
+1 -1
View File
@@ -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,)
nodes/FL_GradGen.py
+2 -2
View File
@@ -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 = {
nodes/FL_OllamaCaptioner.py
+100
View File
@@ -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}'",)
nodes/SaveAndDisplayImage.py
+29
View File
@@ -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,)
}
web/nodes/GradientImageGenerator.js
+2 -2
View File
@@ -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);
}
}
web/nodes/SaveAndDisplayImage.js
+10
View File
@@ -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", "");
}
}
});