fix

.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+config_.py

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

nodes/__init__.py

@@ -0,0 +1,27 @@
+from .channel_select import ChannelSelector
+from .mask_image import MaskImage
+from .save_surreal import SaveJsonToSurreal, SaveTextToSurreal
+from .fetch_surreal import FetchJsonFromSurreal
+from .foreground_mask import ForegroundMask
+from .save_to_s3 import SaveImageToS3
+
+NODE_CLASS_MAPPINGS = {
+    "ChannelSelector": ChannelSelector,
+    "MaskImage": MaskImage,
+    "SaveImageToS3": SaveImageToS3,
+    "SaveJsonToSurreal": SaveJsonToSurreal,
+    "SaveTextToSurreal": SaveTextToSurreal,
+    "FetchJsonFromSurreal": FetchJsonFromSurreal,
+    "ForegroundMask": ForegroundMask,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ChannelSelector":"ChannelSelector",
+    "MaskImage": "MaskImage",
+    "SaveImageToS3": "SaveImageToS3",
+    "SaveJsonToSurreal": "SaveJsonToSurreal",
+    "SaveTextToSurreal": "SaveTextToSurreal",
+    "FetchJsonFromSurreal": "FetchJsonFromSurreal",
+    "ForegroundMask": "ForegroundMask",
+}

nodes/channel_select.py

@@ -0,0 +1,31 @@
+import torch
+
+class ChannelSelector:
+   
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "channel": ("INT", {
+                    "default": 0, 
+                    "min": 0, #Minimum value
+                    "max": 100, #Maximum value
+                    "step": 1, #Slider's step
+                    "display": "number" # Cosmetic only: display as "number" or "slider"
+                })
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "main"
+
+    CATEGORY = "tensorops"
+
+    def main(self, image, channel):
+        # Select the specified channel and add a new dimension at position 0
+        mask = image[channel].unsqueeze(0)
+
+        return (mask,)
+

nodes/config.py

@@ -0,0 +1,9 @@
+SURREAL_URL = ""
+SURREAL_NAMESPACE = ""
+SURREAL_USER = ""
+SURREAL_PASSWORD = ""
+DATALAKE_AWS_ENDPOINT_URL= ""
+DATALAKE_AWS_ACCESS_KEY_ID= ""
+DATALAKE_AWS_SECRET_ACCESS_KEY= ""
+DATALAKE_AWS_REGION= ""
+BUCKET = ""

nodes/fetch_surreal.py

@@ -0,0 +1,30 @@
+from .surreal import surreal_connect
+
+SURREAL_TABLE = "processor"
+
+class FetchJsonFromSurreal:
+   
+    @classmethod    
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "database": ("STRING", {"multiline": False}),
+                "id": ("STRING", {"multiline": False}),
+                "key": ("STRING", {"multiline": False})
+            },
+        }
+
+    RETURN_TYPES = ("JSON",)
+
+    FUNCTION = "main"
+    OUTPUT_NODE = True
+    CATEGORY = "database_ops"
+    
+    def main(self, database: str, id: str, key: str):
+        connection = surreal_connect(database)
+        query = f"SELECT {key} from {SURREAL_TABLE}:`{id}`"
+        results = connection.query(query).result
+        data = results[0][key]
+        print("results", data)
+        return [data]
+

nodes/florence.py

@@ -0,0 +1,459 @@
+import torch
+import torchvision.transforms.functional as F
+import io
+import os
+from typing import List
+import matplotlib
+matplotlib.use('Agg')   
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from PIL import Image, ImageDraw, ImageColor, ImageFont
+import random
+import numpy as np
+import re
+
+#workaround for unnecessary flash_attn requirement
+from unittest.mock import patch
+from transformers.dynamic_module_utils import get_imports
+
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    if not str(filename).endswith("modeling_florence2.py"):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    imports.remove("flash_attn")
+    return imports
+
+
+import comfy.model_management as mm
+from comfy.utils import ProgressBar
+import folder_paths
+
+script_directory = os.path.dirname(os.path.abspath(__file__))
+
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+class DownloadAndLoadFlorence2Model:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "model": (
+                    [ 
+                    'microsoft/Florence-2-base',
+                    'microsoft/Florence-2-base-ft',
+                    'microsoft/Florence-2-large',
+                    'microsoft/Florence-2-large-ft',
+                    'HuggingFaceM4/Florence-2-DocVQA'
+                    ],
+                    {
+                    "default": 'microsoft/Florence-2-base'
+                    }),
+            "precision": ([ 'fp16','bf16','fp32'],
+                    {
+                    "default": 'fp16'
+                    }),
+            "attention": (
+                    [ 'flash_attention_2', 'sdpa', 'eager'],
+                    {
+                    "default": 'sdpa'
+                    }),
+
+            },
+        }
+
+    RETURN_TYPES = ("FL2MODEL",)
+    RETURN_NAMES = ("florence2_model",)
+    FUNCTION = "loadmodel"
+    CATEGORY = "Florence2"
+
+    def loadmodel(self, model, precision, attention):
+        device = mm.get_torch_device()
+        offload_device = mm.unet_offload_device()
+        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
+
+        model_name = model.rsplit('/', 1)[-1]
+        model_path = os.path.join(folder_paths.models_dir, "LLM", model_name)
+        
+        if not os.path.exists(model_path):
+            print(f"Downloading Lumina model to: {model_path}")
+            from huggingface_hub import snapshot_download
+            snapshot_download(repo_id=model,
+                            local_dir=model_path,
+                            local_dir_use_symlinks=False)
+            
+        print(f"using {attention} for attention")
+        with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
+            model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True)
+        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+        
+        florence2_model = {
+            'model': model, 
+            'processor': processor,
+            'dtype': dtype
+            }
+
+        return (florence2_model,)
+
+def calculate_bounding_box(width, height, flat_points) -> List[float]:
+    """
+    Calculate the bounding box for a polygon.
+
+    Args:
+    flat_points (list of int): Flat list of x, y coordinates defining the polygon points.
+
+    Returns:
+    tuple: (min_x, min_y, max_x, max_y) defining the bounding box.
+    """
+    if not flat_points or len(flat_points) % 2 != 0:
+        raise ValueError("The list of points must be non-empty and have an even number of elements")
+
+    x_coords = flat_points[0::2]
+    y_coords = flat_points[1::2]
+
+    min_x = min(x_coords)
+    max_x = max(x_coords)
+    min_y = min(y_coords)
+    max_y = max(y_coords)
+
+    return [min_x / width, min_y / height, max_x / width, max_y / height]
+    
+class Florence2Run:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "florence2_model": ("FL2MODEL", ),
+                "text_input": ("STRING", {"default": "", "multiline": True}),
+                "task": (
+                    [ 
+                    'region_caption',
+                    'dense_region_caption',
+                    'region_proposal',
+                    'caption',
+                    'detailed_caption',
+                    'more_detailed_caption',
+                    'caption_to_phrase_grounding',
+                    'referring_expression_segmentation',
+                    'ocr',
+                    'ocr_with_region',
+                    'docvqa'
+                    ],
+                   ),
+                "fill_mask": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "keep_model_loaded": ("BOOLEAN", {"default": False}),
+                "max_new_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}),
+                "num_beams": ("INT", {"default": 3, "min": 1, "max": 64}),
+                "do_sample": ("BOOLEAN", {"default": True}),
+                "output_mask_select": ("STRING", {"default": ""}),
+            }
+        }
+    
+    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "JSON")
+    RETURN_NAMES =("image", "mask", "caption", "data") 
+    FUNCTION = "encode"
+    CATEGORY = "Florence2"
+
+    def encode(self, image, text_input, florence2_model, task, fill_mask, keep_model_loaded=False, 
+            num_beams=3, max_new_tokens=1024, do_sample=True, output_mask_select=""):
+        device = mm.get_torch_device()
+        _, height, width, _ = image.shape
+        offload_device = mm.unet_offload_device()
+        annotated_image_tensor = None
+        mask_tensor = None
+        processor = florence2_model['processor']
+        model = florence2_model['model']
+        dtype = florence2_model['dtype']
+        model.to(device)
+
+        colormap = ['blue','orange','green','purple','brown','pink','olive','cyan','red',
+                    'lime','indigo','violet','aqua','magenta','gold','tan','skyblue']
+
+        prompts = {
+            'region_caption': '<OD>',
+            'dense_region_caption': '<DENSE_REGION_CAPTION>',
+            'region_proposal': '<REGION_PROPOSAL>',
+            'caption': '<CAPTION>',
+            'detailed_caption': '<DETAILED_CAPTION>',
+            'more_detailed_caption': '<MORE_DETAILED_CAPTION>',
+            'caption_to_phrase_grounding': '<CAPTION_TO_PHRASE_GROUNDING>',
+            'referring_expression_segmentation': '<REFERRING_EXPRESSION_SEGMENTATION>',
+            'ocr': '<OCR>',
+            'ocr_with_region': '<OCR_WITH_REGION>',
+            'docvqa': '<DocVQA>'
+        }
+        task_prompt = prompts.get(task, '<OD>')
+
+        if (task not in ['referring_expression_segmentation', 'caption_to_phrase_grounding', 'docvqa']) and text_input:
+            raise ValueError("Text input (prompt) is only supported for 'referring_expression_segmentation', 'caption_to_phrase_grounding', and 'docvqa'")
+
+        if text_input != "":
+            prompt = task_prompt + " " + text_input
+        else:
+            prompt = task_prompt
+
+        image = image.permute(0, 3, 1, 2)
+        
+        out = []
+        out_masks = []
+        out_results = []
+        out_data = []
+        pbar = ProgressBar(len(image))
+        for img in image:
+            image_pil = F.to_pil_image(img)
+            inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device)
+
+            generated_ids = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                num_beams=num_beams,
+            )
+
+            results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+            print(results)
+            # cleanup the special tokens from the final list
+            if task == 'ocr_with_region':
+                clean_results = str(results)       
+                cleaned_string = re.sub(r'</?s>|<[^>]*>', '\n',  clean_results)
+                clean_results = re.sub(r'\n+', '\n', cleaned_string)
+            else:
+                clean_results = str(results)       
+                clean_results = clean_results.replace('</s>', '')
+                clean_results = clean_results.replace('<s>', '')
+
+            #return single string if only one image for compatibility with nodes that can't handle string lists
+            if len(image) == 1:
+                out_results = clean_results
+            else:
+                out_results.append(clean_results)
+
+            W, H = image_pil.size
+            
+            parsed_answer = processor.post_process_generation(results, task=task_prompt, image_size=(W, H))
+
+            if task == 'region_caption' or task == 'dense_region_caption' or task == 'caption_to_phrase_grounding' or task == 'region_proposal':           
+                fig, ax = plt.subplots(figsize=(W / 100, H / 100), dpi=100)
+                fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
+                ax.imshow(image_pil)
+                bboxes = parsed_answer[task_prompt]['bboxes']
+                labels = parsed_answer[task_prompt]['labels']
+
+                mask_indexes = []
+                # Determine mask indexes outside the loop
+                if output_mask_select != "":
+                    mask_indexes = [n for n in output_mask_select.split(",")]
+                    print(mask_indexes)
+                else:
+                    mask_indexes = [str(i) for i in range(len(bboxes))]
+
+                # Initialize mask_layer only if needed
+                if fill_mask:
+                    mask_layer = Image.new('RGB', image_pil.size, (0, 0, 0))
+                    mask_draw = ImageDraw.Draw(mask_layer)
+
+                for index, (bbox, label) in enumerate(zip(bboxes, labels)):
+                    # Modify the label to include the index
+                    indexed_label = f"{index}.{label}"
+                    
+                    if fill_mask:
+                        if str(index) in mask_indexes:
+                            print("match index:", str(index), "in mask_indexes:", mask_indexes)
+                            mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255))
+                        if label in mask_indexes:
+                            print("match label")
+                            mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255))
+
+                    # Create a Rectangle patch
+                    rect = patches.Rectangle(
+                        (bbox[0], bbox[1]),  # (x,y) - lower left corner
+                        bbox[2] - bbox[0],   # Width
+                        bbox[3] - bbox[1],   # Height
+                        linewidth=1,
+                        edgecolor='r',
+                        facecolor='none',
+                        label=indexed_label
+                    )
+                     # Calculate text width with a rough estimation
+                    text_width = len(label) * 6  # Adjust multiplier based on your font size
+                    text_height = 12  # Adjust based on your font size
+
+                    # Initial text position
+                    text_x = bbox[0]
+                    text_y = bbox[1] - text_height  # Position text above the top-left of the bbox
+
+                    # Adjust text_x if text is going off the left or right edge
+                    if text_x < 0:
+                        text_x = 0
+                    elif text_x + text_width > W:
+                        text_x = W - text_width
+
+                    # Adjust text_y if text is going off the top edge
+                    if text_y < 0:
+                        text_y = bbox[3]  # Move text below the bottom-left of the bbox if it doesn't overlap with bbox
+
+                    # Add the rectangle to the plot
+                    ax.add_patch(rect)
+                    facecolor = random.choice(colormap) if len(image) == 1 else 'red'
+                    # Add the label
+                    plt.text(
+                        text_x,
+                        text_y,
+                        indexed_label,
+                        color='white',
+                        fontsize=12,
+                        bbox=dict(facecolor=facecolor, alpha=0.5)
+                    )
+                if fill_mask:             
+                    mask_tensor = F.to_tensor(mask_layer)
+                    mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
+                    mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
+                    mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
+                    mask_tensor = mask_tensor[:, :, :, 0]
+                    out_masks.append(mask_tensor)           
+
+                # Remove axis and padding around the image
+                ax.axis('off')
+                ax.margins(0,0)
+                ax.get_xaxis().set_major_locator(plt.NullLocator())
+                ax.get_yaxis().set_major_locator(plt.NullLocator())
+                fig.canvas.draw() 
+                buf = io.BytesIO()
+                plt.savefig(buf, format='png', pad_inches=0)
+                buf.seek(0)
+                annotated_image_pil = Image.open(buf)
+
+                annotated_image_tensor = F.to_tensor(annotated_image_pil)
+                out_tensor = annotated_image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
+                out.append(out_tensor)
+
+                
+                pbar.update(1)
+    
+                plt.close(fig)
+
+            elif task == 'referring_expression_segmentation':
+                # Create a new black image
+                mask_image = Image.new('RGB', (W, H), 'black')
+                mask_draw = ImageDraw.Draw(mask_image)
+  
+                predictions = parsed_answer[task_prompt]
+    
+                # Iterate over polygons and labels  
+                for polygons, label in zip(predictions['polygons'], predictions['labels']):
+                    color = random.choice(colormap)
+                    for _polygon in polygons:  
+                        _polygon = np.array(_polygon).reshape(-1, 2)
+                        # Clamp polygon points to image boundaries
+                        _polygon = np.clip(_polygon, [0, 0], [W - 1, H - 1])
+                        if len(_polygon) < 3:  
+                            print('Invalid polygon:', _polygon)
+                            continue  
+                        
+                        _polygon = _polygon.reshape(-1).tolist()
+                        
+                        # Draw the polygon
+                        if fill_mask:
+                            overlay = Image.new('RGBA', image_pil.size, (255, 255, 255, 0))
+                            image_pil = image_pil.convert('RGBA')
+                            draw = ImageDraw.Draw(overlay)
+                            color_with_opacity = ImageColor.getrgb(color) + (180,)
+                            draw.polygon(_polygon, outline=color, fill=color_with_opacity, width=3)
+                            image_pil = Image.alpha_composite(image_pil, overlay)
+                        else:
+                            draw = ImageDraw.Draw(image_pil)
+                            draw.polygon(_polygon, outline=color, width=3)
+
+                        #draw mask
+                        mask_draw.polygon(_polygon, outline="white", fill="white")
+                        
+                image_tensor = F.to_tensor(image_pil)
+                image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() 
+                out.append(image_tensor)
+
+                mask_tensor = F.to_tensor(mask_image)
+                mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
+                mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
+                mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
+                mask_tensor = mask_tensor[:, :, :, 0]
+                out_masks.append(mask_tensor)
+                pbar.update(1)
+
+            elif task == 'ocr_with_region':
+                try:
+                    font = ImageFont.load_default().font_variant(size=24)
+                except:
+                    font = ImageFont.load_default()
+                predictions = parsed_answer[task_prompt]
+                scale = 1
+                draw = ImageDraw.Draw(image_pil)
+                bboxes, labels = predictions['quad_boxes'], predictions['labels']
+                
+                for box, label in zip(bboxes, labels):
+                    bbox = calculate_bounding_box(width, height, box)
+                    out_data.append({"label": label, "polygon": box, "box": bbox})
+                    color = random.choice(colormap)
+                    new_box = (np.array(box) * scale).tolist()
+                    draw.polygon(new_box, width=3, outline=color)
+                    draw.text((new_box[0]+8, new_box[1]+2),
+                                "{}".format(label),
+                                align="right",
+                                font=font,
+                                fill=color)
+                    
+                image_tensor = F.to_tensor(image_pil)
+                image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
+                out.append(image_tensor)
+
+            elif task == 'docvqa':
+                if text_input == "":
+                    raise ValueError("Text input (prompt) is required for 'docvqa'")
+                prompt = "<DocVQA> " + text_input
+
+                inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device)
+                generated_ids = model.generate(
+                    input_ids=inputs["input_ids"],
+                    pixel_values=inputs["pixel_values"],
+                    max_new_tokens=max_new_tokens,
+                    do_sample=do_sample,
+                    num_beams=num_beams,
+                )
+
+                results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+                clean_results = results.replace('</s>', '').replace('<s>', '')
+                
+                if len(image) == 1:
+                    out_results = clean_results
+                else:
+                    out_results.append(clean_results)
+                    
+                out.append(F.to_tensor(image_pil).unsqueeze(0).permute(0, 2, 3, 1).cpu().float())
+
+                pbar.update(1)
+            
+        if len(out) > 0:
+            out_tensor = torch.cat(out, dim=0)
+        else:
+            out_tensor = torch.zeros((1, 64,64, 3), dtype=torch.float32, device="cpu")
+        if len(out_masks) > 0:
+            out_mask_tensor = torch.cat(out_masks, dim=0)
+        else:
+            out_mask_tensor = torch.zeros((1,64,64), dtype=torch.float32, device="cpu")
+
+        if not keep_model_loaded:
+            print("Offloading model...")
+            model.to(offload_device)
+            mm.soft_empty_cache()
+        
+        return (out_tensor, out_mask_tensor, out_results, out_data)
+     
+NODE_CLASS_MAPPINGS = {
+    "DownloadAndLoadFlorence2Model": DownloadAndLoadFlorence2Model,
+    "Florence2Run": Florence2Run,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DownloadAndLoadFlorence2Model": "DownloadAndLoadFlorence2Model",
+    "Florence2Run": "Florence2Run",
+}

nodes/foreground_mask.py

@@ -0,0 +1,192 @@
+from typing import List
+from PIL import Image, ImageDraw
+import numpy as np
+from collections import Counter
+import torchvision.transforms.functional as F
+import json
+
+def calculate_bounding_box(points) -> List[float]:
+    """
+    Calculate the bounding box for a polygon.
+
+    Args:
+    flat_points (list of int): Flat list of x, y coordinates defining the polygon points.
+
+    Returns:
+    tuple: (min_x, min_y, max_x, max_y) defining the bounding box.
+    """
+    if not points or len(points) % 2 != 0:
+        raise ValueError("The list of points must be non-empty and have an even number of elements")
+
+    x_coords = points[0::2]
+    y_coords = points[1::2]
+
+    min_x = min(x_coords)
+    max_x = max(x_coords)
+    min_y = min(y_coords)
+    max_y = max(y_coords)
+
+    return [min_x, min_y, max_x, max_y]
+
+
+def find_mode_color(image: Image.Image):
+    """
+    Identify the most frequent (mode) color in a PIL image.
+    
+    Parameters:
+    image_path (str): The path to the input image.
+    
+    Returns:
+    tuple: The mode color in the image as an (R, G, B) tuple.
+    """
+    # Convert image to RGB mode if it's not already
+    image = image.convert('RGB')
+    
+    # Get the list of pixels
+    pixels = list(image.getdata())
+    
+    # Use Counter to count the frequency of each color
+    counter = Counter(pixels)
+    
+    # Find the most common color
+    mode_color = counter.most_common(1)[0][0]
+    
+    return mode_color
+
+def separate_foreground_background(image):
+    """
+    Separate the Pillow image into foreground and background using the mode color and distance clustering.
+    
+    Parameters:
+    image_path (str): The path to the input image.
+    output_foreground (str): The path to save the foreground image.
+    output_background (str): The path to save the background image.
+    
+    Returns:
+    None
+    """
+    # Convert image to RGBA mode to handle transparency
+    image = image.convert('RGBA')
+    pixels = np.array(image)
+    
+    # Calculate the Euclidean distance of each pixel to the mode color
+    background_color = find_mode_color(image)
+    print("Background color:", background_color)
+    mode_color_array = np.array(background_color)
+    distances = np.linalg.norm(pixels[:, :, :3] - mode_color_array, axis=2)
+    
+    # Determine the threshold distance for clustering
+    threshold_distance = np.mean(distances)
+    
+    print("Threshold distance:", threshold_distance)
+    # Create masks for foreground and background
+    foreground_mask = distances > threshold_distance
+    background_mask = distances <= threshold_distance
+    
+    # Create empty arrays for the new images
+    foreground_image = np.zeros_like(pixels)
+    background_image = np.zeros_like(pixels)
+    
+    # Copy the pixels to the new images based on the masks
+    foreground_image[foreground_mask] = pixels[foreground_mask]
+    background_image[background_mask] = pixels[background_mask]
+    
+    # Find the fg color
+    fg_color = find_mode_color(Image.fromarray(foreground_image, 'RGBA'))
+    
+    # Set foreground pixels with alpha == 255 to black
+    alpha_channel = foreground_image[:, :, 3] == 255
+    foreground_image[alpha_channel, :3] = [255, 255, 255]
+    foreground_image[:, :, 3] = 255
+
+    # Convert back to PIL images
+    foreground_image = Image.fromarray(foreground_image, 'RGBA')
+    background_image = Image.fromarray(background_image, 'RGBA')
+    
+    # Invert Foreground As White
+    # foreground_image = ImageOps.invert(foreground_image.convert("RGB"))
+    
+    return foreground_image, fg_color
+
+def crop_polygon(image, points):
+    """
+    Create a white mask on a black image of size width x height using a list of polygon points.
+
+    Args:
+    points (list of tuples): List of (x, y) tuples defining the polygon points.
+    width (int): Width of the image.
+    height (int): Height of the image.
+
+    Returns:
+    Image: Pillow Image object with the polygon mask.
+    """
+    x_min, y_min, x_max, y_max = calculate_bounding_box(points)
+    image_crop = image.crop((x_min, y_min, x_max, y_max))
+    return image_crop
+    
+def mask_polygon(image, points):
+    """
+    Crop a polygon from a Pillow image.
+
+    Args:
+    image (PIL.Image): The input image.
+    flat_points (list of int): Flat list of x, y coordinates defining the polygon points.
+
+    Returns:
+    PIL.Image: Cropped image of the polygon.
+    """
+    if not points or len(points) % 2 != 0:
+        raise ValueError("The list of points must be non-empty and have an even number of elements")
+
+    # Create a mask
+    mask = Image.new('L', image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    new_box = (np.array(points) * 1.0).tolist()
+    draw.polygon(new_box, fill="white")
+
+    # Apply the mask to the image
+    masked_image = Image.composite(image.convert("RGBA"), mask.convert("RGBA"), mask)
+    return masked_image
+
+
+import torch
+
+class ForegroundMask:
+   
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "json_data": ("JSON",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "main"
+
+    CATEGORY = "tensorops"
+
+    def main(self, image: torch.Tensor, json_data: str):
+        print(json_data)
+        image = image.permute(0, 3, 1, 2)
+        print("Image", image.shape)
+        print("JsonData", json_data, len(json_data))
+        image_pil = F.to_pil_image(image[0])
+        full_image = Image.new("RGBA", image_pil.size, (0, 0, 0, 255))
+        print("ForegroundMask")
+        print("ImageShape", image.shape)
+        print("OutImageShape", full_image.size)
+        for item in json_data:
+            points = item["polygon"]
+            print("polygon", points)
+            masked_image = mask_polygon(image_pil, points)
+            print("NewMaskImageShape", masked_image.size)
+            masked_image_crop = crop_polygon(image_pil, points)
+            fg_image, fg_color = separate_foreground_background(masked_image_crop)
+            x_min, y_min, x_max, y_max = calculate_bounding_box(points)
+            full_image.paste(fg_image, (int(x_min), int(y_min)))
+        out_image = F.to_tensor(full_image)
+        return (out_image,)
+

nodes/mask_image.py

@@ -0,0 +1,28 @@
+import torch
+
+class MaskImage:
+   
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "main"
+
+    CATEGORY = "tensorops"
+
+    def main(self, image: torch.Tensor, mask: torch.Tensor):
+        mask = mask.unsqueeze(-1)
+        new_image = image * mask
+        print("MaskImage")
+        print("ImageShape", image.shape)
+        print("MaskShape", mask.shape)
+        print("NewImageShape", new_image.shape)
+        return (new_image,)
+

nodes/save_surreal.py

@@ -0,0 +1,53 @@
+from .surreal import surreal_connect
+
+SURREAL_TABLE = "processor"
+
+class SaveJsonToSurreal:
+   
+    @classmethod    
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "database": ("STRING", {"multiline": False}),
+                "json": ("JSON",),
+                "id": ("STRING", {"multiline": False}),
+                "key": ("STRING", {"multiline": False})
+            },
+        }
+
+    RETURN_TYPES = ()
+
+    FUNCTION = "main"
+    OUTPUT_NODE = True
+    CATEGORY = "database_ops"
+    
+    def main(self, database: str, id: str, key: str, json: str):
+        connection = surreal_connect(database)
+        query = f"UPDATE {SURREAL_TABLE}:`{id}` CONTENT {{{key}: {json}}};"
+        connection.query(query)
+        return ()
+
+class SaveTextToSurreal:
+   
+    @classmethod    
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "database": ("STRING", {"multiline": False}),
+                "text":  ("STRING",{"forceInput": True}),
+                "id": ("STRING", {"multiline": False}),
+                "key": ("STRING", {"multiline": False})
+            },
+        }
+
+    RETURN_TYPES = ()
+
+    FUNCTION = "main"
+    OUTPUT_NODE = True
+    CATEGORY = "database_ops"
+
+    def main(self, database: str, id: str, key: str, text: str):
+        connection = surreal_connect(database)
+        query = f"UPDATE {SURREAL_TABLE}:`{id}` CONTENT {{{key}: '{text}'}};"
+        connection.query(query)
+        return ()

nodes/save_to_s3.py

@@ -0,0 +1,63 @@
+from .config import DATALAKE_AWS_ACCESS_KEY_ID, DATALAKE_AWS_ENDPOINT_URL, DATALAKE_AWS_SECRET_ACCESS_KEY, BUCKET
+import boto3
+from botocore.client import Config
+import torch
+import PIL
+import PIL.Image
+import io
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+S3_RESOURCE = None
+
+def init_s3():
+    global S3_RESOURCE
+    if not S3_RESOURCE:
+        S3_ENDPOINT_URL = DATALAKE_AWS_ENDPOINT_URL
+        AWS_ACCESS_KEY_ID = DATALAKE_AWS_ACCESS_KEY_ID
+        AWS_SECRET_ACCESS_KEY = DATALAKE_AWS_SECRET_ACCESS_KEY
+        # Initialize the S3 client
+        S3_RESOURCE = boto3.resource(
+            "s3",
+            endpoint_url=S3_ENDPOINT_URL,
+            aws_access_key_id=AWS_ACCESS_KEY_ID,
+            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
+            config=Config(signature_version="s3v4"),
+        )
+       
+def store_image(key: str, image: PIL.Image.Image):
+    global S3_RESOURCE
+    init_s3()
+    try:
+        s3_bucket = S3_RESOURCE.Bucket(BUCKET)
+        file_content = io.BytesIO()
+        image.save(file_content, format="webp")
+        file_content.seek(0)
+        logger.info(f"StoreImage: {key} in {BUCKET}")
+        s3_bucket.put_object(Key=key, Body=file_content)
+    except Exception as e:
+        logger.error(f"StoreContentError {key}: {e}")
+         
+class SaveImageToS3:
+   
+    @classmethod    
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "database": ("STRING", {"multiline": False}),
+                "key": ("STRING", {"multiline": False}),
+                "image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ()
+
+    FUNCTION = "main"
+    OUTPUT_NODE = True
+    CATEGORY = "database_ops"
+    
+    def main(self, database: str, key: str, image: torch.Tensor):
+        store_image(f"{database}/{key}.webp", PIL.Image.fromarray(image.mul(255).byte().permute(1, 2, 0).cpu().numpy()))
+        return ()

nodes/surreal.py

@@ -0,0 +1,7 @@
+from surrealist import Surreal
+from .config import SURREAL_URL, SURREAL_NAMESPACE, SURREAL_USER, SURREAL_PASSWORD
+
+def surreal_connect(database: str):
+    surreal_client = Surreal(SURREAL_URL, namespace=SURREAL_NAMESPACE, database=database, credentials=(SURREAL_USER, SURREAL_PASSWORD), use_http=True, timeout=10)
+    surreal_connection =  surreal_client.connect()
+    return surreal_connection

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "comfyui-tensorops"
+description = "Nodes to perform tensor operations in ComfyUI"
+version = "1.0.0"
+license = "MIT"
+dependencies = ["transformers>=4.38.0"]
+
+[project.urls]
+Repository = "https://github.com/un-seen/comfyui-tensorops"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "shkr"
+DisplayName = "comfyui-tensorops"
+Icon = ""

requirements.txt

@@ -0,0 +1,4 @@
+tqdm
+numpy
+surrealist
+boto3==1.34.86