From 1494bca4591a9bf01c0886bf8d45dd38e9afb674 Mon Sep 17 00:00:00 2001 From: ControlAltAI <6172927+gseth@users.noreply.github.com> Date: Tue, 17 Sep 2024 23:06:52 +0530 Subject: [PATCH] Added Noise Plus Blend. Integer Setting, Choose Upscale Model Nodes --- __init__.py | 15 ++++-- choose_upscale_model_node.py | 30 ++++++++++++ integer_settings_node.py | 28 ++++++++++++ noise_plus_blend_node.py | 89 ++++++++++++++++++++++++++++++++++++ 4 files changed, 159 insertions(+), 3 deletions(-) create mode 100644 choose_upscale_model_node.py create mode 100644 integer_settings_node.py create mode 100644 noise_plus_blend_node.py diff --git a/__init__.py b/__init__.py index 5a5f917..aacfce5 100644 --- a/__init__.py +++ b/__init__.py @@ -4,6 +4,9 @@ from .flux_controlnet_node import FluxControlNetApply from .boolean_basic_node import BooleanBasic from .boolean_reverse_node import BooleanReverse from .get_image_size_ratio_node import GetImageSizeRatio +from .noise_plus_blend_node import NoisePlusBlend +from .integer_settings_node import IntegerSettings +from .choose_upscale_model_node import ChooseUpscaleModel # Import the new node print("Initializing ControlAltAI Nodes") @@ -13,7 +16,10 @@ NODE_CLASS_MAPPINGS = { "FluxControlNetApply": FluxControlNetApply, "BooleanBasic": BooleanBasic, "BooleanReverse": BooleanReverse, - "GetImageSizeRatio": GetImageSizeRatio + "GetImageSizeRatio": GetImageSizeRatio, + "NoisePlusBlend": NoisePlusBlend, + "IntegerSettings": IntegerSettings, + "ChooseUpscaleModel": ChooseUpscaleModel, # Add new node to class mappings } NODE_DISPLAY_NAME_MAPPINGS = { @@ -22,5 +28,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FluxControlNetApply": "Flux ControlNet", "BooleanBasic": "Boolean Basic", "BooleanReverse": "Boolean Reverse", - "GetImageSizeRatio": "Get Image Size & Ratio" -} \ No newline at end of file + "GetImageSizeRatio": "Get Image Size & Ratio", + "NoisePlusBlend": "Noise Plus Blend", + "IntegerSettings": "Integer Settings", + "ChooseUpscaleModel": "Choose Upscale Model", # Add display name mapping for new node +} diff --git a/choose_upscale_model_node.py b/choose_upscale_model_node.py new file mode 100644 index 0000000..dc46743 --- /dev/null +++ b/choose_upscale_model_node.py @@ -0,0 +1,30 @@ +class ChooseUpscaleModel: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "upscale_model_1": ("UPSCALE_MODEL",), + "upscale_model_2": ("UPSCALE_MODEL",), + "use_model_1": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("UPSCALE_MODEL",) + RETURN_NAMES = ("upscale_model",) + FUNCTION = "choose_upscale_model" + + CATEGORY = "ControlAltAI Nodes/Logic" + + def choose_upscale_model(self, upscale_model_1, upscale_model_2, use_model_1): + if use_model_1: + return (upscale_model_1,) + else: + return (upscale_model_2,) + +NODE_CLASS_MAPPINGS = { + "ChooseUpscaleModel": ChooseUpscaleModel, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ChooseUpscaleModel": "Choose Upscale Model", +} diff --git a/integer_settings_node.py b/integer_settings_node.py new file mode 100644 index 0000000..55f22bd --- /dev/null +++ b/integer_settings_node.py @@ -0,0 +1,28 @@ +class IntegerSettings: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "setting": ("BOOLEAN", {"default": False, "label_on": "Enable", "label_off": "Disable"}), + }, + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("setting_value",) + FUNCTION = "integer_settings" + + CATEGORY = "ControlAltAI Nodes/Logic" + + def integer_settings(self, setting): + # Handle the single boolean setting + status = 2 if setting else 1 + return (status,) + + +NODE_CLASS_MAPPINGS = { + "IntegerSettings": IntegerSettings, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IntegerSettings": "Integer Settings", +} diff --git a/noise_plus_blend_node.py b/noise_plus_blend_node.py new file mode 100644 index 0000000..1a0292d --- /dev/null +++ b/noise_plus_blend_node.py @@ -0,0 +1,89 @@ +import numpy as np +from PIL import Image, ImageChops +import torch + +class NoisePlusBlend: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "noise_scale": ("FLOAT", {"default": 0.40, "min": 0.00, "max": 100.00, "step": 0.01}), + "blend_opacity": ("INT", {"default": 20, "min": 0, "max": 100}), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("blended_image_output", "noise_output") + FUNCTION = "noise_plus_blend" + + CATEGORY = "ControlAltAI Nodes/Image" + + def tensor_to_pil(self, tensor_image): + """Converts tensor to a PIL Image""" + tensor_image = tensor_image.squeeze(0) # Remove batch dimension if it exists + pil_image = Image.fromarray((tensor_image.cpu().numpy() * 255).astype(np.uint8)) + return pil_image + + def pil_to_tensor(self, pil_image): + """Converts a PIL image back to a tensor""" + return torch.from_numpy(np.array(pil_image).astype(np.float32) / 255).unsqueeze(0) + + def generate_gaussian_noise(self, width, height, noise_scale=0.05): + """Generates Gaussian noise with a given scale.""" + noise = np.random.normal(128, 128 * noise_scale, (height, width, 3)).astype(np.uint8) + return Image.fromarray(noise) + + def soft_light_blend(self, base_image, noise_image, mask=None, opacity=15): + """Blends noise over the base image using soft light, applying mask if present.""" + # Resize noise to match base image size + noise_image = noise_image.resize(base_image.size) + + base_image = base_image.convert('RGB') + noise_image = noise_image.convert('RGB') + + noise_blended = ImageChops.soft_light(base_image, noise_image) + blended_image = Image.blend(base_image, noise_blended, opacity / 100) + + # Apply mask only if it's provided, valid, and contains more than a single value + if mask is not None: + mask_pil = self.tensor_to_pil(mask).convert('L') + mask_resized = mask_pil.resize(base_image.size) + + # Invert the mask by subtracting from 255 + inverted_mask = ImageChops.invert(mask_resized) + + # Apply the inverted mask to the composite blending + blended_image = Image.composite(base_image, blended_image, inverted_mask) + + return blended_image + + def noise_plus_blend(self, image, noise_scale=0.05, blend_opacity=15, mask=None): + """Main function to generate noise, blend, and return results.""" + # Convert Tensor image to PIL + base_image = self.tensor_to_pil(image) + image_size = base_image.size + + # Generate Gaussian noise with the size of the input image + noise_image = self.generate_gaussian_noise(image_size[0], image_size[1], noise_scale) + + # Blend the noise with the base image using soft light + blended_image = self.soft_light_blend(base_image, noise_image, mask, blend_opacity) + + # Convert the final blended image back to tensor + noise_tensor = self.pil_to_tensor(noise_image) + blended_tensor = self.pil_to_tensor(blended_image) + + # Return both the noise and blended image as tensors + return blended_tensor, noise_tensor + +NODE_CLASS_MAPPINGS = { + "NoisePlusBlend": NoisePlusBlend, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NoisePlusBlend": "Noise Plus Blend", +}