From c462880f3d883df4d6bd831fe7ce59eb616daf3c Mon Sep 17 00:00:00 2001 From: morphicschris Date: Wed, 25 Oct 2023 22:39:20 +0100 Subject: [PATCH] Added some new nodes --- .idea/.gitignore | 8 + .idea/ComfyUI-ImageGlitcher.iml | 10 ++ .../inspectionProfiles/profiles_settings.xml | 6 + .idea/misc.xml | 4 + .idea/modules.xml | 8 + .idea/vcs.xml | 6 + __init__.py | 167 +----------------- nodes/__init__.py | 18 ++ nodes/color_stylizer.py | 100 +++++++++++ nodes/image_glitcher.py | 166 +++++++++++++++++ nodes/local_llm.py | 76 ++++++++ 11 files changed, 404 insertions(+), 165 deletions(-) create mode 100644 .idea/.gitignore create mode 100644 .idea/ComfyUI-ImageGlitcher.iml create mode 100644 .idea/inspectionProfiles/profiles_settings.xml create mode 100644 .idea/misc.xml create mode 100644 .idea/modules.xml create mode 100644 .idea/vcs.xml create mode 100644 nodes/__init__.py create mode 100644 nodes/color_stylizer.py create mode 100644 nodes/image_glitcher.py create mode 100644 nodes/local_llm.py diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..13566b8 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,8 @@ +# Default ignored files +/shelf/ +/workspace.xml +# Editor-based HTTP Client requests +/httpRequests/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/.idea/ComfyUI-ImageGlitcher.iml b/.idea/ComfyUI-ImageGlitcher.iml new file mode 100644 index 0000000..74d515a --- /dev/null +++ b/.idea/ComfyUI-ImageGlitcher.iml @@ -0,0 +1,10 @@ + + + + + + + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100644 index 0000000..105ce2d --- /dev/null +++ b/.idea/inspectionProfiles/profiles_settings.xml @@ -0,0 +1,6 @@ + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..2b2a12e --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,4 @@ + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..5436f4d --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..35eb1dd --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/__init__.py b/__init__.py index f5d872f..d721463 100644 --- a/__init__.py +++ b/__init__.py @@ -1,166 +1,3 @@ -import torch -import random -from torchvision.transforms.functional import to_pil_image, to_tensor -from PIL import ImageEnhance, Image, ImageChops +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -class ImageGlitcher: - """ - Apply a glitch effect on the input image. - """ - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Input: image, glitch_mount, brightness, scanlines - """ - return { - "required": { - "image": ("IMAGE",), - "glitchiness": ("INT", { - "default": 2, - "min": 0, - "max": 10, - "step": 1, - "display": "slider" - }), - "brightness": ("INT", { - "default": 0, - "min": 0, - "max": 10, - "step": 1, - "display": "slider" - }), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "scanlines": (["enable", "disable"],) - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "applyGlitch" - OUTPUT_NODE = False - CATEGORY = "Effects" - - def applyGlitch(self, image, glitchiness, brightness, scanlines, seed): - # Since the glitchImage method isn't provided, - # I'm just adding a placeholder for the logic. - # In a real scenario, you'd call glitchImage and pass the required parameters. - glitched_image = self.glitchImage(image, glitchiness, brightness, scanlines == "enable", seed) - return (glitched_image,) - - - def glitchImage(self, tensor, glitch_amount, brightness_amount, use_scanlines, seed): - random.seed(seed) - - # Ensure the tensor is of shape (B, H, W, C) - tensor = tensor.squeeze(0).permute(0, 1, 2).cpu().numpy() - img = Image.fromarray((tensor * 255).astype('uint8')) - - iw, ih = img.size - max_offset = int(glitch_amount * glitch_amount / 100 * iw) - - # Create output image and input image copies - output_img = img.copy() - input_img = img.copy() - - # Randomly offset slices horizontally - for _ in range(glitch_amount * 2): - startY = random.randint(0, ih) - chunk_height = random.randint(1, ih // 4) - chunk_height = min(chunk_height, ih - startY) - offset = random.randint(-max_offset, max_offset) - - if offset == 0: - continue - - # Left shift and wrap-around - if offset < 0: - output_img.paste(input_img.crop((0, startY, iw + offset, startY + chunk_height)), (0, startY)) - output_img.paste(input_img.crop((iw + offset, startY, iw, startY + chunk_height)), (0, startY)) - else: - # Right shift and wrap-around - output_img.paste(input_img.crop((offset, startY, iw, startY + chunk_height)), (0, startY)) - output_img.paste(input_img.crop((0, startY, offset, startY + chunk_height)), (iw - offset, startY)) - - # Color Offset - channel_to_offset = self.get_random_channel() - offset_x = random.randint(-glitch_amount * 2, glitch_amount * 2) - offset_y = random.randint(-glitch_amount * 2, glitch_amount * 2) - r, g, b = img.split() - r2, g2, b2 = output_img.split() - - if channel_to_offset == 'R': - r = ImageChops.offset(r, offset_x, offset_y) - g = g2 - b = b2 - elif channel_to_offset == 'G': - r = r2 - g = ImageChops.offset(g, offset_x, offset_y) - b = b2 - elif channel_to_offset == 'B': - r = r2 - g = g2 - b = ImageChops.offset(b, offset_x, offset_y) - output_img = Image.merge("RGB", (r, g, b)) -# output_img = self.blend_single_channel(input_img, output_img, channel_to_offset, 0) - - # Brightness - enhancer = ImageEnhance.Brightness(output_img) - output_img = enhancer.enhance(1 + brightness_amount / 10) - - # Add Scanlines - if use_scanlines: - for i in range(ih): - if i % 2 == 0: - line = Image.new("RGB", (iw, 1), (0, 0, 0)) - output_img.paste(line, (0, i)) - - # Convert back to tensor and maintain (B, H, W, C) format - glitched_tensor = to_tensor(output_img).unsqueeze(0).permute(0, 2, 3, 1) - - return glitched_tensor - - def get_random_channel(self): - r = random.random() - if r < 0.33: - return 'G' - elif r < 0.66: - return 'R' - else: - return 'B' - - def blend_single_channel(self, img1, img2, channel_to_blend, alpha=0.5): - # Ensure both images are in RGB mode - img1 = img1.convert("RGB") - img2 = img2.convert("RGB") - - # Split both images into their R, G, and B channels - r1, g1, b1 = img1.split() - r2, g2, b2 = img2.split() - - # Perform blending operation on the chosen channel - if channel_to_blend == "R": - blended_r = Image.blend(r1, r2, alpha) - result = Image.merge("RGB", (blended_r, g1, b1)) - elif channel_to_blend == "G": - blended_g = Image.blend(g1, g2, alpha) - result = Image.merge("RGB", (r1, blended_g, b1)) - elif channel_to_blend == "B": - blended_b = Image.blend(b1, b2, alpha) - result = Image.merge("RGB", (r1, g1, blended_b)) - else: - raise ValueError("Invalid channel. Choose from 'R', 'G', or 'B'.") - - return result - -# Append to the NODE_CLASS_MAPPINGS dictionary -NODE_CLASS_MAPPINGS = { - "ImageGlitcher": ImageGlitcher -} - -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = { - "ImageGlitcher": "Image Glitcher" -} \ No newline at end of file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..687aa75 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,18 @@ +from .image_glitcher import ImageGlitcher +from .color_stylizer import ColorStylizer +from .local_llm import QueryLocalLLM + +NODE_CLASS_MAPPINGS = { + "ImageGlitcher": ImageGlitcher, + "ColorStylizer": ColorStylizer, + "QueryLocalLLM": QueryLocalLLM, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageGlitcher": "Image Glitcher", + "ColorStylizer": "Color Stylizer", + "QueryLocalLLM": "Query Local LLM", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/color_stylizer.py b/nodes/color_stylizer.py new file mode 100644 index 0000000..4cb0a8b --- /dev/null +++ b/nodes/color_stylizer.py @@ -0,0 +1,100 @@ +import torch +import cv2 +import numpy as np +from torchvision.transforms.functional import to_tensor + +class ColorStylizer: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "target_r": ("INT", { + "default": 255, + "min": 0, + "max": 255, + "step": 1, + "display": "slider" + }), + "target_g": ("INT", { + "default": 255, + "min": 0, + "max": 255, + "step": 1, + "display": "slider" + }), + "target_b": ("INT", { + "default": 255, + "min": 0, + "max": 255, + "step": 1, + "display": "slider" + }), + "falloff": ("FLOAT", { + "default": 30.0, + "min": 0.0, + "max": 100.0, + "step": 1.0, + "display": "slider" + }), + "gain": ("FLOAT", { + "default": 1.5, + "min": 0.0, + "max": 10.0, + "step": 0.5, + "display": "slider" + }) + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "stylize" + OUTPUT_NODE = False + CATEGORY = "Effects" + + def stylize(self, image, target_r, target_g, target_b, falloff, gain): + target_color = (target_b, target_g, target_r) + image = image.squeeze(0) + image = image.mul(255).byte().numpy() + image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + gray_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + gray_img = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2BGR) + falloff_mask = self.create_falloff_mask(image, target_color, falloff) + image_amplified = image.copy() + image_amplified[:, :, 2] = np.clip(image_amplified[:, :, 2] * gain, 0, 255).astype(np.uint8) + stylized_img = (image_amplified * falloff_mask + gray_img * (1 - falloff_mask)).astype(np.uint8) + stylized_img = cv2.cvtColor(stylized_img, cv2.COLOR_BGR2RGB) + stylized_img_tensor = to_tensor(stylized_img).float() + stylized_img_tensor = stylized_img_tensor.permute(1, 2, 0).unsqueeze(0) + return (stylized_img_tensor,) + + def create_falloff_mask(self, img, target_color, falloff): + target_color = np.array(target_color, dtype=np.uint8) + print("img shape:", img.shape) + print("img dtype:", img.dtype) + print("target_color shape:", target_color.shape) + print("target_color dtype:", target_color.dtype) + target_color = np.full_like(img, target_color) + print(111) + diff = cv2.absdiff(img, target_color) + print(222) + diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY) + print(333) + _, mask = cv2.threshold(diff, falloff, 255, cv2.THRESH_BINARY_INV) + mask = cv2.GaussianBlur(mask, (0, 0), falloff / 2) + mask = mask / 255.0 + mask = mask.reshape(*mask.shape, 1) + return mask + +# Append to the NODE_CLASS_MAPPINGS dictionary +NODE_CLASS_MAPPINGS = { + "ColorStylizer": ColorStylizer +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ColorStylizer": "Color Stylizer" +} diff --git a/nodes/image_glitcher.py b/nodes/image_glitcher.py new file mode 100644 index 0000000..f5d872f --- /dev/null +++ b/nodes/image_glitcher.py @@ -0,0 +1,166 @@ +import torch +import random +from torchvision.transforms.functional import to_pil_image, to_tensor +from PIL import ImageEnhance, Image, ImageChops + +class ImageGlitcher: + """ + Apply a glitch effect on the input image. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Input: image, glitch_mount, brightness, scanlines + """ + return { + "required": { + "image": ("IMAGE",), + "glitchiness": ("INT", { + "default": 2, + "min": 0, + "max": 10, + "step": 1, + "display": "slider" + }), + "brightness": ("INT", { + "default": 0, + "min": 0, + "max": 10, + "step": 1, + "display": "slider" + }), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "scanlines": (["enable", "disable"],) + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "applyGlitch" + OUTPUT_NODE = False + CATEGORY = "Effects" + + def applyGlitch(self, image, glitchiness, brightness, scanlines, seed): + # Since the glitchImage method isn't provided, + # I'm just adding a placeholder for the logic. + # In a real scenario, you'd call glitchImage and pass the required parameters. + glitched_image = self.glitchImage(image, glitchiness, brightness, scanlines == "enable", seed) + return (glitched_image,) + + + def glitchImage(self, tensor, glitch_amount, brightness_amount, use_scanlines, seed): + random.seed(seed) + + # Ensure the tensor is of shape (B, H, W, C) + tensor = tensor.squeeze(0).permute(0, 1, 2).cpu().numpy() + img = Image.fromarray((tensor * 255).astype('uint8')) + + iw, ih = img.size + max_offset = int(glitch_amount * glitch_amount / 100 * iw) + + # Create output image and input image copies + output_img = img.copy() + input_img = img.copy() + + # Randomly offset slices horizontally + for _ in range(glitch_amount * 2): + startY = random.randint(0, ih) + chunk_height = random.randint(1, ih // 4) + chunk_height = min(chunk_height, ih - startY) + offset = random.randint(-max_offset, max_offset) + + if offset == 0: + continue + + # Left shift and wrap-around + if offset < 0: + output_img.paste(input_img.crop((0, startY, iw + offset, startY + chunk_height)), (0, startY)) + output_img.paste(input_img.crop((iw + offset, startY, iw, startY + chunk_height)), (0, startY)) + else: + # Right shift and wrap-around + output_img.paste(input_img.crop((offset, startY, iw, startY + chunk_height)), (0, startY)) + output_img.paste(input_img.crop((0, startY, offset, startY + chunk_height)), (iw - offset, startY)) + + # Color Offset + channel_to_offset = self.get_random_channel() + offset_x = random.randint(-glitch_amount * 2, glitch_amount * 2) + offset_y = random.randint(-glitch_amount * 2, glitch_amount * 2) + r, g, b = img.split() + r2, g2, b2 = output_img.split() + + if channel_to_offset == 'R': + r = ImageChops.offset(r, offset_x, offset_y) + g = g2 + b = b2 + elif channel_to_offset == 'G': + r = r2 + g = ImageChops.offset(g, offset_x, offset_y) + b = b2 + elif channel_to_offset == 'B': + r = r2 + g = g2 + b = ImageChops.offset(b, offset_x, offset_y) + output_img = Image.merge("RGB", (r, g, b)) +# output_img = self.blend_single_channel(input_img, output_img, channel_to_offset, 0) + + # Brightness + enhancer = ImageEnhance.Brightness(output_img) + output_img = enhancer.enhance(1 + brightness_amount / 10) + + # Add Scanlines + if use_scanlines: + for i in range(ih): + if i % 2 == 0: + line = Image.new("RGB", (iw, 1), (0, 0, 0)) + output_img.paste(line, (0, i)) + + # Convert back to tensor and maintain (B, H, W, C) format + glitched_tensor = to_tensor(output_img).unsqueeze(0).permute(0, 2, 3, 1) + + return glitched_tensor + + def get_random_channel(self): + r = random.random() + if r < 0.33: + return 'G' + elif r < 0.66: + return 'R' + else: + return 'B' + + def blend_single_channel(self, img1, img2, channel_to_blend, alpha=0.5): + # Ensure both images are in RGB mode + img1 = img1.convert("RGB") + img2 = img2.convert("RGB") + + # Split both images into their R, G, and B channels + r1, g1, b1 = img1.split() + r2, g2, b2 = img2.split() + + # Perform blending operation on the chosen channel + if channel_to_blend == "R": + blended_r = Image.blend(r1, r2, alpha) + result = Image.merge("RGB", (blended_r, g1, b1)) + elif channel_to_blend == "G": + blended_g = Image.blend(g1, g2, alpha) + result = Image.merge("RGB", (r1, blended_g, b1)) + elif channel_to_blend == "B": + blended_b = Image.blend(b1, b2, alpha) + result = Image.merge("RGB", (r1, g1, blended_b)) + else: + raise ValueError("Invalid channel. Choose from 'R', 'G', or 'B'.") + + return result + +# Append to the NODE_CLASS_MAPPINGS dictionary +NODE_CLASS_MAPPINGS = { + "ImageGlitcher": ImageGlitcher +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageGlitcher": "Image Glitcher" +} \ No newline at end of file diff --git a/nodes/local_llm.py b/nodes/local_llm.py new file mode 100644 index 0000000..bf01f98 --- /dev/null +++ b/nodes/local_llm.py @@ -0,0 +1,76 @@ +import requests +import json + +class QueryLocalLLM: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "prompt": ("STRING", { "multiline": True, "default": "" }), + "url": ("STRING", { "multiline": False, "default": "http://127.0.0.1:5000/api/v1/generate" }), + "context_length": ("INT", { "default": 2048, "min": 512, "max": 4096, "display": "slider" }), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}) + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("generated_text",) + FUNCTION = "generateText" + OUTPUT_NODE = False + CATEGORY = "Effects" + + def generateText(self, prompt, url, context_length, seed): + description = self.call_api(prompt, url, context_length, seed) + return (description,) + + def call_api(self, prompt_text, url, context_length, seed): + payload = { + "prompt": f"{prompt_text}\nAssistant: ", + "use_story": False, + "use_memory": False, + "use_authors_note": False, + "use_world_info": False, + "max_context_length": context_length, + "max_length": 300, + "rep_pen": 1.1, + "rep_pen_range": 600, + "rep_pen_slope": 0, + "temperature": 1, + "tfs": 1, + "top_a": 0, + "top_k": 0, + "top_p": 0.95, + "typical": 1, + "sampler_order": [6, 0, 1, 2, 3, 4, 5], + "singleline": False, + "seed": seed + } + + # Sending the POST request + response = requests.post(url, json=payload) + + # Checking for successful response + if response.status_code == 200: + result_json = response.json() + + generatedText = result_json["results"][0]["text"] + + print("Response: " + generatedText) + + return generatedText + else: + print(f"Error {response.status_code}: {response.text}") + return None + +# Append to the NODE_CLASS_MAPPINGS dictionary +NODE_CLASS_MAPPINGS = { + "QueryLocalLLM": QueryLocalLLM +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "QueryLocalLLM": "Query Local LLM" +} \ No newline at end of file