301 lines
11 KiB
Python
301 lines
11 KiB
Python
from nodes import interrupt_processing
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import torch
|
|
import numpy as np
|
|
|
|
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
|
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
loop_interation = 0
|
|
reset_global = False
|
|
class BookToolsPromptSelector:
|
|
"""
|
|
Selects and concatenates values from a dictionary based on input indices and ranges.
|
|
The input string contains 1-based indices which are converted to dictionary keys.
|
|
Assumes the dictionary values are strings which will be concatenated and returned as a single string.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
"""
|
|
Describes expected input types:
|
|
- 'dictionary': a dictionary where keys will be selected based on 'selected_indexes'.
|
|
- 'selected_indexes': a int listing keys by index, separated by commas, with optional ranges.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"dictionary": ("DICTIONARY", ),
|
|
"selected_indexes": ("INT", {
|
|
"multiline": False,
|
|
"default": "1,2,3"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", )
|
|
FUNCTION = "main"
|
|
CATEGORY = "selector"
|
|
OUTPUT_NODE = False
|
|
|
|
def main(self, dictionary: dict, selected_indexes: int):
|
|
if not isinstance(dictionary, dict):
|
|
raise ValueError("Conditioning must be a dictionary.")
|
|
|
|
if not isinstance(selected_indexes, int):
|
|
raise ValueError("selected_indexes must be an integer.")
|
|
|
|
result_string = str(dictionary.get(selected_indexes, ''))
|
|
|
|
return (result_string,)
|
|
|
|
class BookToolsPromptSchedule:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"text": ("STRING", {"multiline": True, "default":"\"first page description\",\n\"second page description\""}),
|
|
"before_text": ("STRING", {"multiline": True,}),
|
|
"after_text": ("STRING", {"multiline": True,}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("DICTIONARY",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "prompt"
|
|
OUTPUT_NODE = False
|
|
|
|
def main(self, text, before_text, after_text):
|
|
"""
|
|
Splits the text on commas, applies before and after text to each segment,
|
|
and returns a dictionary with each formatted text indexed.
|
|
|
|
Parameters:
|
|
- text (str): Comma-separated text entries to be formatted.
|
|
- before_text (str): Text to prepend to each entry.
|
|
- after_text (str): Text to append to each entry.
|
|
|
|
Returns:
|
|
- dict: Dictionary with indexed formatted texts.
|
|
"""
|
|
segments = text.split('",')
|
|
formatted_texts = {}
|
|
|
|
for index, segment in enumerate(segments, start=1):
|
|
cleaned_segment = segment.replace('"', '').strip()
|
|
parts = [part for part in [before_text.strip(), cleaned_segment, after_text.strip()] if part]
|
|
formatted_text = ', '.join(parts)
|
|
formatted_texts[str(index)] = formatted_text
|
|
return (formatted_texts,)
|
|
|
|
# based on https://civitai.com/models/26836/comfyui-loopback-nodes but take any parametrs and have reset option
|
|
class BookToolsLoop:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"reset": ("BOOLEAN", {"default": False}),}}
|
|
|
|
RETURN_TYPES = ("LOOP", "INT",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "loopback"
|
|
RETURN_NAMES = ("LOOP", "Iteration",)
|
|
|
|
def run(self, reset):
|
|
global loop_interation, reset_global
|
|
if (reset == True):
|
|
loop_interation = 1
|
|
reset_global = True
|
|
else:
|
|
reset_global = False
|
|
return (self, loop_interation)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, reset):
|
|
global loop_interation, reset_global
|
|
loop_interation +=1
|
|
if (reset == True):
|
|
loop_interation = 1
|
|
reset_global = True
|
|
else:
|
|
reset_global = False
|
|
return loop_interation
|
|
|
|
class BookToolsLoopStart:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"first_loop": (any,),
|
|
"loop": ("LOOP",)
|
|
}
|
|
}
|
|
|
|
FUNCTION = "run"
|
|
CATEGORY = "loopback"
|
|
RETURN_TYPES = (any,)
|
|
RETURN_NAMES = ("*",)
|
|
|
|
def run(self, first_loop, loop):
|
|
global reset_global
|
|
if reset_global == True:
|
|
return (first_loop,)
|
|
if hasattr(loop, 'next'):
|
|
return (loop.next,)
|
|
return (first_loop,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, first_loop, loop):
|
|
global reset_global
|
|
if reset_global == True:
|
|
return (first_loop,)
|
|
if hasattr(loop, 'next'):
|
|
return id(loop.next)
|
|
return float("NaN")
|
|
|
|
class BookToolsLoopEnd:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "send_to_next_loop": (any, ), "loop": ("LOOP",) }}
|
|
FUNCTION = "run"
|
|
CATEGORY = "loopback"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
|
|
def run(self, send_to_next_loop, loop):
|
|
loop.next = send_to_next_loop
|
|
return ()
|
|
|
|
class BookToolsEndQueue:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "boolean": ("BOOLEAN", {"default": False}), }}
|
|
FUNCTION = "main"
|
|
CATEGORY = "logic"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
|
|
def main(self, boolean):
|
|
if (boolean==True):
|
|
interrupt_processing()
|
|
return ()
|
|
|
|
class BookToolsImageTextOverlay:
|
|
def __init__(self, device="cpu"):
|
|
self.device = device
|
|
_alignments = ["left", "right", "center"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"text": ("STRING", {"multiline": True, "default": "Hello"}),
|
|
"textbox_width": ("INT", {"default": 200, "min": 1}),
|
|
"textbox_height": ("INT", {"default": 200, "min": 1}),
|
|
"max_font_size": ("INT", {"default": 30, "min": 1, "max": 256, "step": 1}),
|
|
"font": ("STRING", {"default": "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"}),
|
|
"alignment": (cls._alignments, {"default": "center"}),
|
|
"color": ("STRING", {"default": "#000000"}),
|
|
"start_x": ("INT", {"default": 0}),
|
|
"start_y": ("INT", {"default": 0}),
|
|
"padding": ("INT", {"default": 50}),
|
|
"line_height": ("INT", {"default": 20, "min": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "add_text_overlay"
|
|
CATEGORY = "image/text"
|
|
|
|
def wrap_text_and_calculate_height(self, text, font, max_width, line_height):
|
|
wrapped_lines = []
|
|
# Split the input text by newline characters to respect manual line breaks
|
|
paragraphs = text.split('\n')
|
|
|
|
for paragraph in paragraphs:
|
|
words = paragraph.split()
|
|
current_line = words[0] if words else ''
|
|
|
|
for word in words[1:]:
|
|
# Test if adding a new word exceeds the max width
|
|
test_line = current_line + ' ' + word if current_line else word
|
|
test_line_bbox = font.getbbox(test_line)
|
|
w = test_line_bbox[2] - test_line_bbox[0] # Right - Left for width
|
|
if w <= max_width:
|
|
current_line = test_line
|
|
else:
|
|
# If the current line plus the new word exceeds max width, wrap it
|
|
wrapped_lines.append(current_line)
|
|
current_line = word
|
|
|
|
# Don't forget to add the last line of the paragraph
|
|
wrapped_lines.append(current_line)
|
|
|
|
# Calculate the total height considering the custom line height
|
|
total_height = len(wrapped_lines) * line_height
|
|
|
|
wrapped_text = '\n'.join(wrapped_lines)
|
|
return wrapped_text, total_height
|
|
|
|
|
|
def add_text_overlay(self, image, text, textbox_width, textbox_height, max_font_size, font, alignment, color, start_x, start_y, padding, line_height):
|
|
image_tensor = image
|
|
image_np = image_tensor.cpu().numpy()
|
|
image_pil = Image.fromarray((image_np.squeeze(0) * 255).astype(np.uint8))
|
|
color_rgb = tuple(int(color.lstrip("#")[i:i+2], 16) for i in (0, 2, 4))
|
|
|
|
effective_textbox_width = textbox_width - 2 * padding # Adjust for padding
|
|
effective_textbox_height = textbox_height - 2 * padding
|
|
|
|
font_size = max_font_size
|
|
while font_size >= 1:
|
|
loaded_font = ImageFont.truetype(font, font_size)
|
|
wrapped_text, total_text_height = self.wrap_text_and_calculate_height(text, loaded_font, effective_textbox_width, line_height)
|
|
|
|
if total_text_height <= effective_textbox_height:
|
|
draw = ImageDraw.Draw(image_pil)
|
|
lines = wrapped_text.split('\n')
|
|
y = start_y + padding + (effective_textbox_height - total_text_height) // 2
|
|
|
|
for line in lines:
|
|
line_bbox = loaded_font.getbbox(line)
|
|
line_width = line_bbox[2] - line_bbox[0]
|
|
|
|
if alignment == "left":
|
|
x = start_x + padding
|
|
elif alignment == "right":
|
|
x = start_x + effective_textbox_width - line_width + padding
|
|
elif alignment == "center":
|
|
x = start_x + padding + (effective_textbox_width - line_width) // 2
|
|
|
|
draw.text((x, y), line, fill=color_rgb, font=loaded_font)
|
|
y += line_height # Use custom line height for spacing
|
|
|
|
break # Break the loop if text fits within the specified dimensions
|
|
|
|
font_size -= 1 # Decrease font size and try again
|
|
|
|
image_tensor_out = torch.tensor(np.array(image_pil).astype(np.float32) / 255.0)
|
|
image_tensor_out = torch.unsqueeze(image_tensor_out, 0)
|
|
return (image_tensor_out,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"BTPromptSelector": BookToolsPromptSelector,
|
|
"BTPromptSchedule": BookToolsPromptSchedule,
|
|
"Loop": BookToolsLoop,
|
|
"LoopStart": BookToolsLoopStart,
|
|
"LoopEnd": BookToolsLoopEnd,
|
|
"EndQueue": BookToolsEndQueue,
|
|
"ImageTextOverlay": BookToolsImageTextOverlay,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"BTPromptSelector": "[Book Tools] Prompt Selector",
|
|
"BTPromptSchedule": "[Book Tools] Prompt Batch Schedule",
|
|
"Loop": "[Book Tools] Loop",
|
|
"LoopStart": "[Book Tools] Loop Start",
|
|
"LoopEnd": "[Book Tools] Loop End",
|
|
"EndQueue": "[Book Tools] End Queue",
|
|
"ImageTextOverlay": "[Book Tools] Image Text Overlay"
|
|
} |