commit 418c0f787d709ed442560e41f75f98770e153d9e Author: or Date: Sun Sep 29 17:44:47 2024 +0300 initial commit diff --git a/Readme.md b/Readme.md new file mode 100644 index 0000000..71f8fe5 --- /dev/null +++ b/Readme.md @@ -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: + +Original image + +An illustration of the workflow: + + Eraser workflow example Original image + +## 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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..215ce66 --- /dev/null +++ b/__init__.py @@ -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", +} diff --git a/bria_api_node.py b/bria_api_node.py new file mode 100644 index 0000000..1e53469 --- /dev/null +++ b/bria_api_node.py @@ -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) diff --git a/images/eraser_workflow.png b/images/eraser_workflow.png new file mode 100644 index 0000000..c16e31e Binary files /dev/null and b/images/eraser_workflow.png differ diff --git a/images/eraser_workflow_diagram.jpg b/images/eraser_workflow_diagram.jpg new file mode 100644 index 0000000..db6aa99 Binary files /dev/null and b/images/eraser_workflow_diagram.jpg differ diff --git a/images/original_image.jpg b/images/original_image.jpg new file mode 100644 index 0000000..1bec2c1 Binary files /dev/null and b/images/original_image.jpg differ diff --git a/workflow.json b/workflow.json new file mode 100644 index 0000000..07a883b --- /dev/null +++ b/workflow.json @@ -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 +} \ No newline at end of file