initial commit

2024-09-29 17:44:47 +03:00
commit 418c0f787d
7 changed files with 369 additions and 0 deletions
@@ -0,0 +1,43 @@
+# BRIA ComfyUI API Nodes
+
+## Overview
+This repository contains custom nodes for ComfyUI that allow access to BRIA's API endpoints.
+
+To use the nodes in the workflow, you need a valid BRIA API token. You can get one [here](https://bria.ai/api/)
+
+You can load the workflow, which includes all available nodes, by importing the [workflow.json](workflow.json) file in this repo. 
+
+You can also download the following image and import it to comfyui:
+
+<img src="./images/eraser_workflow.png" alt="Original image" width="500"/>
+
+An illustration of the workflow:
+
+ <img src="./images/eraser_workflow_diagram.jpg" alt="Eraser workflow example" width="650"/> <img src="./images/original_image.jpg" alt="Original image" width="150"/>
+
+## Available Nodes
+
+### Eraser
+The **Eraser** node allows users to remove specific objects or areas from an image by providing a mask.
+
+This functionality is powered by BRIA's ControlNet inpainting, available on [this model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Inpainting) on Hugging Face.
+
+## Installation
+There are two methods to install the BRIA ComfyUI API nodes:
+
+### Method 1: Using ComfyUI's Custom Node Manager
+1. Open ComfyUI.
+2. Navigate to the [**Custom Node Manager**](https://github.com/ltdrdata/ComfyUI-Manager).
+3. Click on 'Install Missing Nodes' or search for BRIA API and install the node from the manager.
+
+### Method 2: Git Clone
+1. Navigate to the `custom_nodes` directory of your ComfyUI installation:
+   ```bash
+   cd path_to_comfyui/custom_nodes
+   ```
+2. Clone this repository:
+   ```bash
+   git clone https://github.com/your-repo-link/ComfyUI-BRIA-API.git
+   ```
+
+3. Restart ComfyUI and load the workflows.
@@ -0,0 +1,11 @@
+from .bria_api_node import EraserNode
+
+# Map the node class to a name used internally by ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "BriaEraser": EraserNode,  # Return the class, not an instance
+}
+
+# Map the node display name to the one shown in the ComfyUI node interface
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BriaEraser": "Bria Eraser",
+}
@@ -0,0 +1,112 @@
+import numpy as np
+import requests
+from PIL import Image
+import io
+import base64
+from torchvision.transforms import ToPILImage, ToTensor
+import torch
+
+# Base class for shared functionality between both nodes
+class BriaAPINode:
+    def __init__(self, api_url):
+        self.api_url = api_url
+    
+    def preprocess_image(self, image):
+        if isinstance(image, torch.Tensor):
+            # Print image shape for debugging
+            if image.dim() == 4:  # (batch_size, height, width, channels)
+                image = image.squeeze(0)  # Remove the batch dimension (1)
+                # Convert to PIL after permuting to (height, width, channels)
+                image = ToPILImage()(image.permute(2, 0, 1))  # (height, width, channels)
+            else:
+                print("Unexpected image dimensions. Expected 4D tensor.")
+        return image
+
+    def preprocess_mask(self, mask):
+        if isinstance(mask, torch.Tensor):
+            # Print mask shape for debugging
+            if mask.dim() == 3:  # (batch_size, height, width)
+                mask = mask.squeeze(0)  # Remove the batch dimension (1)
+                # Convert to PIL (grayscale mask)
+                mask = ToPILImage()(mask)  # No permute needed for grayscale
+            else:
+                print("Unexpected mask dimensions. Expected 3D tensor.")
+        return mask
+    
+    def image_to_base64(self, pil_image):
+        # Convert a PIL image to a base64-encoded string
+        buffered = io.BytesIO()
+        pil_image.save(buffered, format="PNG")  # Save the image to the buffer in PNG format
+        buffered.seek(0)  # Rewind the buffer to the beginning
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    def process_request(self, image, mask, api_key):
+        if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
+            raise Exception("Please insert a valid API key.")
+
+        # Check if image and mask are tensors, if so, convert to NumPy arrays
+        if isinstance(image, torch.Tensor):
+            image = self.preprocess_image(image)
+        if isinstance(mask, torch.Tensor):
+            mask = self.preprocess_mask(mask)
+
+        # Convert the image and mask directly to Base64 strings
+        image_base64 = self.image_to_base64(image)
+        mask_base64 = self.image_to_base64(mask)
+
+        # Prepare the API request payload
+        payload = {
+            "file": f"{image_base64}",
+            "mask_file": f"{mask_base64}"
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "api_token": f"{api_key}"
+        }
+
+        try:
+            response = requests.post(self.api_url, json=payload, headers=headers)
+            # Check for successful response
+            if response.status_code == 200:
+                print('response is 200')
+                # Process the output image from API response
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result_url'])
+                result_image = Image.open(io.BytesIO(image_response.content))
+                result_image = result_image.convert("RGBA")
+                result_image = np.array(result_image).astype(np.float32) / 255.0
+                result_image = torch.from_numpy(result_image)[None,]
+                # image_tensor = image_tensor = ToTensor()(output_image)
+                # image_tensor = image_tensor.permute(1, 2, 0) / 255.0  # Shape now becomes [1, 2200, 1548, 3]
+                # print(f"output tensor shape is: {image_tensor.shape}")
+                return (result_image,)
+            else:
+                raise Exception(f"Error: API request failed with status code {response.status_code}")
+
+        except Exception as e:
+            raise Exception(f"{e}")
+
+# Eraser Node
+class EraserNode(BriaAPINode):
+    @staticmethod
+    def INPUT_TYPES():
+        return {
+            "required": {
+                "image": ("IMAGE",),  # Input image from another node
+                "mask": ("MASK",),  # Binary mask input
+                "api_key": ("STRING", {"default": "BRIA_API_TOKEN"})  # API Key input with a default value
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image",)
+    CATEGORY = "API Nodes"
+    FUNCTION = "execute"  # This is the method that will be executed
+
+    def __init__(self):
+        super().__init__("https://engine.prod.bria-api.com/v1/eraser")  # Eraser API URL
+
+    # Define the execute method as expected by ComfyUI
+    def execute(self, image, mask, api_key):
+        return self.process_request(image, mask, api_key)
@@ -0,0 +1,203 @@
+{
+  "last_node_id": 28,
+  "last_link_id": 42,
+  "nodes": [
+    {
+      "id": 15,
+      "type": "Note",
+      "pos": {
+        "0": 1021,
+        "1": 280
+      },
+      "size": {
+        "0": 311.8914794921875,
+        "1": 153.69827270507812
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [],
+      "properties": {},
+      "widgets_values": [
+        "The default BRIA API key for ComfyUI (BRIA_ComfyUI_Key) offers 10,000 API calls for the entire community. \n\nGet your own token at:\nhttps://bria.ai/api/"
+      ],
+      "color": "#432",
+      "bgcolor": "#653"
+    },
+    {
+      "id": 13,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 1410,
+        "1": 160
+      },
+      "size": {
+        "0": 474.7605895996094,
+        "1": 303.117919921875
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 42
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    },
+    {
+      "id": 21,
+      "type": "LoadImage",
+      "pos": {
+        "0": 477,
+        "1": 156
+      },
+      "size": {
+        "0": 408.4602355957031,
+        "1": 333.19830322265625
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            40
+          ],
+          "slot_index": 0,
+          "shape": 3
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": [
+            41
+          ],
+          "slot_index": 1,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "clipspace/clipspace-mask-82245.69999998808.png [input]",
+        "image"
+      ]
+    },
+    {
+      "id": 28,
+      "type": "BriaEraser",
+      "pos": {
+        "0": 1022,
+        "1": 159
+      },
+      "size": {
+        "0": 315,
+        "1": 78
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 40
+        },
+        {
+          "name": "mask",
+          "type": "MASK",
+          "link": 41
+        }
+      ],
+      "outputs": [
+        {
+          "name": "output_image",
+          "type": "IMAGE",
+          "links": [
+            42
+          ],
+          "slot_index": 0,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "BriaEraser"
+      },
+      "widgets_values": [
+        "BRIA_ComfyUI_Key"
+      ]
+    },
+    {
+      "id": 14,
+      "type": "Note",
+      "pos": {
+        "0": 483,
+        "1": 39
+      },
+      "size": [
+        396.80859375,
+        61.8046875
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [],
+      "properties": {},
+      "widgets_values": [
+        "Right click, and choose \"Open in Mask Editor\" to draw a mask of areas you want animated more. "
+      ],
+      "color": "#432",
+      "bgcolor": "#653"
+    }
+  ],
+  "links": [
+    [
+      40,
+      21,
+      0,
+      28,
+      0,
+      "IMAGE"
+    ],
+    [
+      41,
+      21,
+      1,
+      28,
+      1,
+      "MASK"
+    ],
+    [
+      42,
+      28,
+      0,
+      13,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        -293.75,
+        167.65625
+      ]
+    }
+  },
+  "version": 0.4
+}