Initial commit

2024-05-31 12:15:30 +02:00
parent b970a9a8eb
commit d3ed002cf1
4 changed files with 306 additions and 0 deletions
@@ -0,0 +1,133 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import math
+
+import comfy.utils
+import comfy.model_management
+
+
+
+
+class Blend:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image1": ("IMAGE",),
+                "image2": ("IMAGE",),
+                "blend_factor": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light", "difference", "add"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "blend_images"
+
+    CATEGORY = "Badman"
+
+    def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str):
+        image2 = image2.to(image1.device)
+        if image1.shape != image2.shape:
+            image2 = image2.permute(0, 3, 1, 2)
+            image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')
+            image2 = image2.permute(0, 2, 3, 1)
+
+        blended_image = self.blend_mode(image1, image2, blend_mode)
+        blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
+        blended_image = torch.clamp(blended_image, 0, 1)
+        return (blended_image,)
+
+    def blend_mode(self, img1, img2, mode):
+        if mode == "normal":
+            return img2
+        elif mode == "multiply":
+            return img1 * img2
+        elif mode == "add":
+            return img1 + img2
+        elif mode == "screen":
+            return 1 - (1 - img1) * (1 - img2)
+        elif mode == "overlay":
+            return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
+        elif mode == "soft_light":
+            return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1))
+        elif mode == "difference":
+            return img1 - img2
+        else:
+            raise ValueError(f"Unsupported blend mode: {mode}")
+
+    def g(self, x):
+        return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
+
+class HexGenerator:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "r": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "g": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "b": ("FLOAT", {
+                    "default": 0.5,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01
+                }),
+                "grayscale": ("BOOLEAN",),
+                
+            },
+        }
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "linear_rgb_to_int"
+
+    CATEGORY = "Badman"
+
+    def linear_rgb_to_int(self,r, g, b, grayscale=False):
+        """
+        Converts linear RGB values to an integer color code.
+
+        Parameters:
+        r (float): Red value (0.0-1.0)
+        g (float): Green value (0.0-1.0)
+        b (float): Blue value (0.0-1.0)
+        grayscale (bool): If True, use the r value for all RGB components
+
+        Returns:
+        int: Integer color code in the format 0xRRGGBB
+        """
+    
+        if grayscale:
+            g = b = r
+        
+        # Convert float values to int (0-255)
+        r_int = int(round(r * 255))
+        g_int = int(round(g * 255))
+        b_int = int(round(b * 255))
+        
+        # Combine into a single integer
+        color_int = (r_int << 16) + (g_int << 8) + b_int
+        return (color_int,)
+
+
@@ -0,0 +1,45 @@
+
+
+class BadmanIntUtil:
+    def __init__(self):
+        pass
+    
+
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "Int1":  ("INT", {
+                    "default": 0,
+                    "min": 0}),
+                "Int2":  ("INT", {
+                    "default": 0,
+                    "min": 0}),
+                "math_function": (["add", "multiply", "sub", "divide"],),
+            },
+        }
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "process_int"
+
+    CATEGORY = "Badman"
+
+    def process_int(self, Int1: int, Int2: int, math_function: str):
+      
+        processed_int = self.math_mode(Int1, Int2, math_function)
+    
+        return (processed_int,)
+
+    def math_mode(self, Int1, Int2, mode):
+        if mode == "add":
+            return Int1+Int2
+        elif mode == "multiply":
+            return Int1 * Int2
+        elif mode == "sub":
+            return Int1 - Int2
+        elif mode == "divide":
+            return Int1/Int1
+        else:
+            raise ValueError(f"Unsupported math mode: {mode}")
+
@@ -0,0 +1,104 @@
+import os
+
+class BadmanPrint:
+
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {"value": ("STRING", {"default": ""})},
+        }
+    
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "log_input"
+    CATEGORY = "Badman"
+
+    def log_input(self, value):
+        print(value,)
+        return {}
+    
+class ConcatString:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+             "required": {
+                "string_field_0": ("STRING", {
+                    "multiline": False,
+                    "default": "Hello"
+                }),
+                "string_field_1": ("STRING", {
+                    "multiline": False,
+                    "default": "World"
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("STRING",)
+    FUNCTION = "execute"
+    CATEGORY = "Badman"
+
+    def execute(self, string_field_0, string_field_1):
+        return (string_field_0 + string_field_1,)
+
+class BadmanString:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {"value": ("STRING", {"default": ""})},
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("STRING",)
+    FUNCTION = "execute"
+    CATEGORY = "Badman"
+
+    def execute(self, value):
+        return (value,)
+    
+class BadmanIOConfigurator:
+    # Get the directory of the current file
+    current_file_directory = os.path.dirname(os.path.abspath(__file__))
+
+    # Get the parent directory
+    Output_directory = os.path.dirname(os.path.dirname(current_file_directory))+"\\output\\"
+ 
+    
+    @classmethod
+    def INPUT_TYPES(s):
+         return {
+            "required": {
+                "BasePath": ("STRING", {"default": s.Output_directory}),
+                "ShotName": ("STRING", {"default": ""}),
+                },
+        }
+    
+    RETURN_TYPES = ("STRING","STRING","STRING","STRING","STRING","STRING","STRING","STRING","STRING",)
+    RETURN_NAMES = ("ImageOutputPath","VideoOutputPath","ImageUpresOutputPath","VideoUpresOutputPath","ImageInputPath","DepthInputPath","PoseInputPath","BodyMaskInputPath","FaceMaskInputPath",)
+    FUNCTION = "execute"
+    CATEGORY = "Badman"
+
+
+    def execute(self, BasePath,ShotName):
+        suffixOut = BasePath+ShotName
+        imagesOutPath = suffixOut+"\Images\image"
+        videoOutPath = suffixOut+"\Video"
+        imagesUpresOutPath = suffixOut+"\ImagesUpres\image"
+        videoUpresOutPath = suffixOut+ "\VideoUpres"
+        framesInPath = BasePath+"ExtractedFrames\\"+ShotName+"\\"
+        depthInPath = BasePath+"ExtractedFrames\\"+ShotName+"\\"
+        poseInPath = BasePath+"ExtractedPose\\"+ShotName+"\\"
+        BodyMaskInPath = BasePath+"Masks\\"+ShotName+"\\Body\\"
+        FaceMaskInPath = BasePath+"Masks\\"+ShotName+"\\Face\\"
+       
+        return (imagesOutPath,videoOutPath,imagesUpresOutPath,videoUpresOutPath,framesInPath,depthInPath,poseInPath,BodyMaskInPath,FaceMaskInPath)
+       
+
+
@@ -0,0 +1,24 @@
+from .BadmanStrings import *  
+from .BadmanImage import * 
+from .BadmanNumbers import * 
+
+
+NODE_CLASS_MAPPINGS = {
+    "Badman_Blend": Blend,
+    "Badman_HexGenerator": HexGenerator,
+    "Badman_String": BadmanString,
+    "Badman_Concat_String": ConcatString,
+    "Badman_Print": BadmanPrint,
+    "BadmanIO" : BadmanIOConfigurator,
+    "BadmanIntUtil" : BadmanIntUtil,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Badman_Blend": "ImageBlend(Badman)",
+    "Badman_HexGenerator": "HexGenerator(Badman)",
+    "Badman_String": "String (Badman)",
+    "Badman_Concat_String": "Concat String (Badman)",
+    "Badman_Print": "Print (Badman)",
+    "Badman_IO": "IO Config (Badman)",
+    "BadmanIntUtil": "Int Math (Badman)",
+}