WAI-3976-feedback-fixes

WAI-3976

.github/workflows/publish.yml

@@ -7,17 +7,19 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
-    # if this is a forked repository. Skipping the workflow.
-    if: github.event.repository.fork == false 
+    if: ${{ github.repository_owner == 'Bria-AI' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           ## Add your own personal access token to your Github Repository secrets and reference it here.
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -1 +1,2 @@
-*.pyc
+*.pyc
+.idea

Readme.md

@@ -13,7 +13,7 @@ An API token is required to use the nodes in your workflows. Get started quickly
   <img src="https://img.shields.io/badge/GET%20YOUR%20TOKEN-1000%20Free%20Calls-blue?style=flat-square" alt="Get Your Token" height="20">
 </a>.
 
-For direct API Endpoint use, look for the endpoint in our of our API partners like: [**fal.ai**](https://fal.ai/models?keywords=bria). 
+for direct API endpoint use, you can find our APIs through partners like [**fal.ai**](https://fal.ai/models?keywords=bria).  
 For source code and weigths access, go to our [**Hugging Face**](https://huggingface.co/briaai) space.
 
 To load a workflow, import the compatible workflow.json files from this [folder](workflows).  
@@ -45,7 +45,8 @@ These nodes use pre-trained tailored models to generate images that faithfully r
 | Node                   | Description                                                        |
 |------------------------|--------------------------------------------------------------------|
 | **Tailored Gen**       | Generates images using a trained tailored model, reproducing specific visual IP elements or guidelines. Use the Tailored Model Info node to load the model's default settings. |
-| **Tailored Model Info** | Retrieves the default settings and prompt prefix of a trained tailored model, which can be used to configure the Tailored Gen node. |
+| **Tailored Model Info**| Retrieves the default settings and prompt prefix of a trained tailored model, which can be used to configure the Tailored Gen node. |
+| **Restyle Portrait**   | Transforms the style of a portrait while preserving the person's facial features. |
 
 ## Image Editing Nodes
 These nodes modify specific parts of images, enabling adjustments while maintaining the integrity of the rest of the image.

__init__.py

@@ -1,5 +1,5 @@
 from .nodes import (EraserNode, GenFillNode, ImageExpansionNode, ReplaceBgNode, RmbgNode, RemoveForegroundNode, ShotByTextNode, ShotByImageNode, TailoredGenNode,
-                     TailoredModelInfoNode, Text2ImageBaseNode, Text2ImageFastNode, Text2ImageHDNode,
+                     TailoredModelInfoNode, Text2ImageBaseNode, Text2ImageFastNode, Text2ImageHDNode, TailoredPortraitNode,
                      ReimagineNode)
 # Map the node class to a name used internally by ComfyUI
 NODE_CLASS_MAPPINGS = {
@@ -13,6 +13,7 @@ NODE_CLASS_MAPPINGS = {
     "ShotByImageNode": ShotByImageNode,
     "BriaTailoredGen": TailoredGenNode,
     "TailoredModelInfoNode": TailoredModelInfoNode,
+    "TailoredPortraitNode": TailoredPortraitNode,
     "Text2ImageBaseNode": Text2ImageBaseNode,
     "Text2ImageFastNode": Text2ImageFastNode,
     "Text2ImageHDNode": Text2ImageHDNode,
@@ -30,6 +31,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ShotByImageNode": "Bria Shot By Image",
     "BriaTailoredGen": "Bria Tailored Gen",
     "TailoredModelInfoNode": "Bria Tailored Model Info",
+    "TailoredPortraitNode": "Bria Restyle Portrait",
     "Text2ImageBaseNode": "Bria Text2Image Base",
     "Text2ImageFastNode": "Bria Text2Image Fast",
     "Text2ImageHDNode": "Bria Text2Image HD",

nodes/__init__.py

@@ -8,6 +8,7 @@ from .shot_by_text_node import ShotByTextNode
 from .shot_by_image_node import ShotByImageNode
 from .tailored_gen_node import TailoredGenNode
 from .tailored_model_info_node import TailoredModelInfoNode
+from .tailored_portrait_node import TailoredPortraitNode
 from .text_2_image_base_node import Text2ImageBaseNode
 from .text_2_image_fast_node import Text2ImageFastNode
 from .text_2_image_hd_node import Text2ImageHDNode

nodes/common.py

@@ -5,6 +5,7 @@ import torch
 import base64
 from torchvision.transforms import ToPILImage
 import requests
+import time
 
 def postprocess_image(image):
     result_image = Image.open(io.BytesIO(image))
@@ -44,7 +45,7 @@ def preprocess_mask(mask):
     return mask
 
 
-def process_request(api_url, image, mask, api_key):
+def process_request(api_url, image, mask, api_key, visual_input_content_moderation, visual_output_content_moderation):
     if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
         raise Exception("Please insert a valid API key.")
 
@@ -58,10 +59,12 @@ def process_request(api_url, image, mask, api_key):
     image_base64 = image_to_base64(image)
     mask_base64 = image_to_base64(mask)
 
-    # Prepare the API request payload
+    # Prepare the API request payload for v2 API
     payload = {
-        "file": f"{image_base64}",
-        "mask_file": f"{mask_base64}"
+        "image": image_base64,
+        "mask": mask_base64,
+        "visual_input_content_moderation":visual_input_content_moderation,
+        "visual_output_content_moderation":visual_output_content_moderation
     }
 
     headers = {
@@ -70,13 +73,23 @@ def process_request(api_url, image, mask, api_key):
     }
 
     try:
-        response = requests.post(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 = requests.post(api_url, json=payload, headers=headers) 
+        if response.status_code == 200 or response.status_code == 202:
+            print('Initial request successful, polling for completion...')
             response_dict = response.json()
-            image_response = requests.get(response_dict['result_url'])
+            status_url = response_dict.get('status_url')
+            request_id = response_dict.get('request_id')
+            
+            if not status_url:
+                raise Exception("No status_url returned from API")
+            
+            print(f"Request ID: {request_id}, Status URL: {status_url}")
+            
+            final_response = poll_status_until_completed(status_url, api_key)
+            result_image_url = final_response['result']['image_url']
+            
+            # Download and process the result image
+            image_response = requests.get(result_image_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
@@ -84,9 +97,51 @@ def process_request(api_url, image, mask, api_key):
             # 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,)
+            return (result_image,)  
         else:
-            raise Exception(f"Error: API request failed with status code {response.status_code}")
+            raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}")
 
     except Exception as e:
         raise Exception(f"{e}")
+
+
+def poll_status_until_completed(status_url, api_key, timeout=360, check_interval=2):
+    """
+    Poll a status URL until the status is COMPLETED or timeout is reached.
+    
+    Args:
+        status_url (str): The status URL to poll
+        api_key (str): API token for authentication
+        timeout (int): Maximum time to wait in seconds (default: 360)
+        check_interval (int): Time between checks in seconds (default: 2)
+    
+    Returns:
+        dict: The final response containing the result
+    
+    Raises:
+        Exception: If timeout is reached or API request fails
+    """
+    start_time = time.time()
+    headers = {"api_token": api_key}
+    
+    while time.time() - start_time < timeout:
+        try:
+            response = requests.get(status_url, headers=headers)            
+            if response.status_code == 200 or response.status_code == 202:
+                response_dict = response.json()
+                status = response_dict.get("status", "").upper()
+                
+                if status == "COMPLETED":
+                    return response_dict
+                elif status == "ERROR":
+                    raise Exception(f"Request failed: {response_dict}")
+                else:
+                    print(f"Status: {status}, waiting...")
+                    time.sleep(check_interval)
+            else:
+                raise Exception(f"Status check failed with status code {response.status_code}")
+                
+        except requests.exceptions.RequestException as e:
+            raise Exception(f"Error checking status: {e}")
+    
+    raise Exception(f"Timeout reached after {timeout} seconds")

nodes/eraser_node.py

@@ -8,6 +8,10 @@ class EraserNode():
                 "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
+            },
+            "optional": {
+                "visual_input_content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}), 
             }
         }
 
@@ -17,9 +21,9 @@ class EraserNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/eraser"  # Eraser API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/erase"  # Eraser API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, mask, api_key):
-        return process_request(self.api_url, image, mask, api_key)
+    def execute(self, image, mask, api_key, visual_input_content_moderation, visual_output_content_moderation):
+        return process_request(self.api_url, image, mask, api_key, visual_input_content_moderation, visual_output_content_moderation)
     

nodes/generative_fill_node.py

@@ -4,7 +4,7 @@ from PIL import Image
 import io
 import torch
 
-from .common import image_to_base64, preprocess_image, preprocess_mask
+from .common import preprocess_image, preprocess_mask, image_to_base64, poll_status_until_completed
 
 
 class GenFillNode():
@@ -18,7 +18,12 @@ class GenFillNode():
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
             "optional": {
-                "seed": ("INT", {"default": 123456})
+                "seed": ("INT", {"default": 123456}),
+                "prompt_content_moderation": ("BOOLEAN", {"default": True}), 
+                "visual_input_content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}), 
+
+
             }
         }
 
@@ -28,10 +33,10 @@ class GenFillNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/gen_fill"  # Eraser API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/gen_fill"
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, mask, prompt, api_key, seed):
+    def execute(self, image, mask, prompt, api_key, seed, prompt_content_moderation, visual_input_content_moderation, visual_output_content_moderation):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -47,12 +52,14 @@ class GenFillNode():
 
         # Prepare the API request payload
         payload = {
-            "file": f"{image_base64}",
-            "mask_file": f"{mask_base64}",
+            "image": image_base64,
+            "mask": mask_base64,
             "prompt": prompt,
             "negative_prompt": "blurry",
-            "sync": True,
             "seed": seed,
+            "prompt_content_moderation":prompt_content_moderation,
+            "visual_input_content_moderation":visual_input_content_moderation,
+            "visual_output_content_moderation":visual_output_content_moderation
         }
 
         headers = {
@@ -61,20 +68,30 @@ class GenFillNode():
         }
 
         try:
+            # Send initial request to get status URL
             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
+            
+            if response.status_code == 200 or response.status_code == 202:
+                print('Initial genfill request successful, polling for completion...')
                 response_dict = response.json()
-                image_response = requests.get(response_dict['urls'][0])
+                status_url = response_dict.get('status_url')
+                request_id = response_dict.get('request_id')
+                
+                if not status_url:
+                    raise Exception("No status_url returned from API")
+                
+                print(f"Request ID: {request_id}, Status URL: {status_url}")
+                
+                final_response = poll_status_until_completed(status_url, api_key)                
+                result_image_url = final_response['result']['image_url']
+                image_response = requests.get(result_image_url)
                 result_image = Image.open(io.BytesIO(image_response.content))
                 result_image = result_image.convert("RGB")
                 result_image = np.array(result_image).astype(np.float32) / 255.0
                 result_image = torch.from_numpy(result_image)[None,]
                 return (result_image,)
             else:
-                raise Exception(f"Error: API request failed with status code {response.status_code}")
+                raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}")
 
         except Exception as e:
             raise Exception(f"{e}")

nodes/image_expansion_node.py

@@ -4,7 +4,7 @@ from PIL import Image
 import io
 import torch
 
-from .common import image_to_base64, preprocess_image
+from .common import image_to_base64, preprocess_image, poll_status_until_completed
 
 
 class ImageExpansionNode():
@@ -13,17 +13,21 @@ class ImageExpansionNode():
         return {
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
-                "original_image_size": ("STRING",),
-                "original_image_location": ("STRING",),
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
             
             "optional": {
+                "original_image_size": ("STRING",),
+                "original_image_location": ("STRING",),
                 "canvas_size": ("STRING", {"default": "1000, 1000"}),
+                "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9","None"], {"default": "None"}),
                 "prompt": ("STRING", {"default": ""}),
                 "seed": ("INT", {"default": 681794}),
                 "negative_prompt": ("STRING", {"default": "Ugly, mutated"}),
-                "content_moderation": ("BOOLEAN", {"default": False}), 
+                "prompt_content_moderation": ("BOOLEAN", {"default": False}),
+                "preserve_alpha": ("BOOLEAN", {"default": True}),
+                "visual_input_content_moderation": ("BOOLEAN", {"default": False}),
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}),
             } 
         }
 
@@ -33,27 +37,28 @@ class ImageExpansionNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/image_expansion"  # Image Expansion API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/expand"  # Image Expansion API URL
 
     # Define the execute method as expected by ComfyUI
     def execute(self, image, 
                 original_image_size,
                 original_image_location,
                 canvas_size,
+                aspect_ratio,
                 prompt,
                 seed,
                 negative_prompt,
-                content_moderation,
+                prompt_content_moderation,
+                preserve_alpha,
+                visual_input_content_moderation,
+                visual_output_content_moderation,
                 api_key):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
+        original_image_size = [int(x.strip()) for x in original_image_size.split(",")] if original_image_size else ()
+        original_image_location = [int(x.strip()) for x in original_image_location.split(",")] if original_image_location else ()
+        canvas_size = [int(x.strip()) for x in canvas_size.split(",")] if canvas_size else ()
         
-        original_image_size = [int(x.strip()) for x in original_image_size.split(",")]
-        original_image_location = [int(x.strip()) for x in original_image_location.split(",")]
-        canvas_size = [int(x.strip()) for x in canvas_size.split(",")]
-        
-        if prompt == "":
-            prompt = None
         if negative_prompt == "":
             negative_prompt = " " # hack to avoid error in triton which expects non-empty string
 
@@ -63,18 +68,32 @@ class ImageExpansionNode():
 
         # Convert the image directly to Base64 string
         image_base64 = image_to_base64(image)
-
-        # Prepare the API request payload
-        payload = {
-            "file": f"{image_base64}",
-            "original_image_size": original_image_size,
-            "original_image_location": original_image_location,
-            "canvas_size": canvas_size,
+        if aspect_ratio and aspect_ratio != "None":
+            payload = {
+            "image": image_base64,
+            "aspect_ratio": aspect_ratio,
             "prompt": prompt,
             "negative_prompt": negative_prompt,
             "seed": seed,
-            "content_moderation": content_moderation
+            "prompt_content_moderation": prompt_content_moderation,
+            "preserve_alpha": preserve_alpha,
+            "visual_input_content_moderation": visual_input_content_moderation,
+            "visual_output_content_moderation": visual_output_content_moderation
         }
+        else:
+            payload = {
+                "image": image_base64,
+                "original_image_size": original_image_size,
+                "original_image_location": original_image_location,
+                "canvas_size": canvas_size,
+                "prompt": prompt,
+                "negative_prompt": negative_prompt,
+                "seed": seed,
+                "prompt_content_moderation": prompt_content_moderation,
+                "preserve_alpha": preserve_alpha,
+                "visual_input_content_moderation": visual_input_content_moderation,
+                "visual_output_content_moderation": visual_output_content_moderation
+            }
 
         headers = {
             "Content-Type": "application/json",
@@ -83,19 +102,34 @@ class ImageExpansionNode():
 
         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
+            
+            if response.status_code == 200 or response.status_code == 202:
+                print('Initial image expansion request successful, polling for completion...')
                 response_dict = response.json()
-                image_response = requests.get(response_dict['result_url'])
+                status_url = response_dict.get('status_url')
+                request_id = response_dict.get('request_id')
+                
+                if not status_url:
+                    raise Exception("No status_url returned from API")
+                
+                print(f"Request ID: {request_id}, Status URL: {status_url}")
+                
+                # Poll status URL until completion
+                final_response = poll_status_until_completed(status_url, api_key)
+                
+                # Get the result image URL
+                result_image_url = final_response['result']['image_url']
+                
+                # Download and process the result image
+                image_response = requests.get(result_image_url)
                 result_image = Image.open(io.BytesIO(image_response.content))
                 result_image = result_image.convert("RGB")
                 result_image = np.array(result_image).astype(np.float32) / 255.0
                 result_image = torch.from_numpy(result_image)[None,]
+                
                 return (result_image,)
             else:
-                raise Exception(f"Error: API request failed with status code {response.status_code}")
+                raise Exception(f"Error: API request failed with status code {response.status_code}: {response.text}")
 
         except Exception as e:
             raise Exception(f"{e}")

nodes/remove_foreground_node.py

@@ -4,7 +4,7 @@ from PIL import Image
 import io
 import torch
 
-from .common import preprocess_image, image_to_base64
+from .common import preprocess_image, image_to_base64, poll_status_until_completed
 
 
 class RemoveForegroundNode():
@@ -16,7 +16,9 @@ class RemoveForegroundNode():
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
             "optional": {
-                "content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_input_content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}), 
+                "preserve_alpha": ("BOOLEAN", {"default": True}), 
             }
         }
 
@@ -26,10 +28,10 @@ class RemoveForegroundNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.internal.prod.bria-api.com/v1/erase_foreground"  # remove foreground API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/erase_foreground"  # remove foreground API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, content_moderation, api_key):
+    def execute(self, image, visual_input_content_moderation, visual_output_content_moderation, preserve_alpha, api_key):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -44,7 +46,12 @@ class RemoveForegroundNode():
         
 #         files=[('file',('temp_img.jpeg', open(temp_img_path, 'rb'),'image/jpeg'))
 #   ]
-        payload = {"file": image_to_base64(image), "content_moderation": content_moderation}
+        payload = {
+            "image": image_to_base64(image),
+            "visual_input_content_moderation": visual_input_content_moderation,
+            "visual_output_content_moderation":visual_output_content_moderation,
+            "preserve_alpha": preserve_alpha
+        }
 
         headers = {
             "Content-Type": "application/json",
@@ -53,18 +60,31 @@ class RemoveForegroundNode():
 
         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
+
+            if response.status_code == 200 or response.status_code == 202:
+                print('Initial request successful, polling for completion...')
                 response_dict = response.json()
-                image_response = requests.get(response_dict['result_url'])
+                status_url = response_dict.get('status_url')
+                request_id = response_dict.get('request_id')
+                
+                if not status_url:
+                    raise Exception("No status_url returned from API")
+                
+                print(f"Request ID: {request_id}, Status URL: {status_url}")
+                
+                # Poll status URL until completion
+                final_response = poll_status_until_completed(status_url, api_key)
+                
+                # Get the result image URL
+                result_image_url = final_response['result']['image_url']
+                
+                image_response = requests.get(result_image_url)
                 result_image = Image.open(io.BytesIO(image_response.content))
                 result_image = np.array(result_image).astype(np.float32) / 255.0
                 result_image = torch.from_numpy(result_image)[None,]
                 return (result_image,)
             else:
-                raise Exception(f"Error: API request failed with status code {response.status_code}")
+                raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}")
 
         except Exception as e:
             raise Exception(f"{e}")

nodes/replace_bg_node.py

@@ -4,7 +4,7 @@ from PIL import Image
 import io
 import torch
 
-from .common import image_to_base64, preprocess_image, preprocess_mask
+from .common import image_to_base64, preprocess_image, preprocess_mask, poll_status_until_completed
 
 
 class ReplaceBgNode():
@@ -16,16 +16,17 @@ class ReplaceBgNode():
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
             "optional": {
-                "fast": ("BOOLEAN", {"default": True}), 
-                "bg_prompt": ("STRING",),
-                "ref_image": ("IMAGE",), # Input ref image from another node
+                "mode": (["base", "fast", "high_control"], {"default": "base"}),
+                "prompt": ("STRING",),
+                "ref_images": ("IMAGE",),
                 "refine_prompt": ("BOOLEAN", {"default": True}), 
-                "enhance_ref_image": ("BOOLEAN", {"default": True}), 
+                "enhance_ref_images": ("BOOLEAN", {"default": True}), 
                 "original_quality": ("BOOLEAN", {"default": False}), 
-                "force_rmbg": ("BOOLEAN", {"default": False}), 
                 "negative_prompt": ("STRING", {"default": None}), 
                 "seed": ("INT", {"default": 681794}),
-                "content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}), 
+                "prompt_content_moderation": ("BOOLEAN", {"default": False}),
+                "force_background_detection": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -35,20 +36,21 @@ class ReplaceBgNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/background/replace"  # Replace BG API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/replace_background"  # Replace BG API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, fast,
+    def execute(self, image, mode,
                 refine_prompt,
-                enhance_ref_image,
                 original_quality,
-                force_rmbg,
                 negative_prompt,
                 seed,
                 api_key,
-                content_moderation,
-                bg_prompt=None,
-                ref_image=None,):
+                visual_output_content_moderation,
+                prompt_content_moderation,
+                enhance_ref_images,
+                force_background_detection,
+                prompt=None,
+                ref_images=None,):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -56,28 +58,29 @@ class ReplaceBgNode():
         if isinstance(image, torch.Tensor):
             image = preprocess_image(image)
 
-        # Convert the image and mask directly to Base64 strings
+        # Convert the image to Base64 string
         image_base64 = image_to_base64(image)
-        ref_image_file = None # initialization, will be updated if it is supplied
-        if ref_image is not None:
-            ref_image = preprocess_image(ref_image)
-            ref_image_file = image_to_base64(ref_image)
+        
+        if ref_images is not None:
+            ref_images = preprocess_image(ref_images)
+            ref_images = [image_to_base64(ref_images)]
+        else:
+            ref_images=[]
 
-        # Prepare the API request payload
+        # Prepare the API request payload for v2 API
         payload = {
-            "file": f"{image_base64}",
-            "fast": fast,
-            "bg_prompt": bg_prompt,
-            "ref_image_file": ref_image_file,
+            "image": image_base64,
+            "mode": mode,
+            "prompt": prompt,
+            "ref_images":ref_images,
             "refine_prompt": refine_prompt,
-            "enhance_ref_image": enhance_ref_image,
             "original_quality": original_quality,
-            "force_rmbg": force_rmbg,
             "negative_prompt": negative_prompt,
             "seed": seed,
-            "sync": True,
-            "num_results": 1,
-            "content_moderation": content_moderation            
+            "prompt_content_moderation": prompt_content_moderation,
+            "visual_output_content_moderation":visual_output_content_moderation,
+            "enhance_ref_images":enhance_ref_images,
+            "force_background_detection": force_background_detection        
         }   
 
         headers = {
@@ -87,19 +90,34 @@ class ReplaceBgNode():
 
         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
+            
+            if response.status_code == 200 or response.status_code == 202:
+                print('Initial replace background request successful, polling for completion...')
                 response_dict = response.json()
-                image_response = requests.get(response_dict['result'][0][0]) # first indexing for batched, second for url
+                status_url = response_dict.get('status_url')
+                request_id = response_dict.get('request_id')
+                
+                if not status_url:
+                    raise Exception("No status_url returned from API")
+                
+                print(f"Request ID: {request_id}, Status URL: {status_url}")
+                
+                # Poll status URL until completion
+                final_response = poll_status_until_completed(status_url, api_key)
+                
+                # Get the result image URL
+                result_image_url = final_response['result']['image_url']
+                
+                # Download and process the result image
+                image_response = requests.get(result_image_url)
                 result_image = Image.open(io.BytesIO(image_response.content))
                 result_image = result_image.convert("RGB")
                 result_image = np.array(result_image).astype(np.float32) / 255.0
                 result_image = torch.from_numpy(result_image)[None,]
+                
                 return (result_image,)
             else:
-                raise Exception(f"Error: API request failed with status code {response.status_code}")
+                raise Exception(f"Error: API request failed with status code {response.status_code}{response.text}")
 
         except Exception as e:
             raise Exception(f"{e}")

nodes/rmbg_node.py

@@ -4,8 +4,7 @@ from PIL import Image
 import io
 import torch
 
-from .common import preprocess_image
-from io import BytesIO
+from .common import preprocess_image, image_to_base64, poll_status_until_completed
 
 class RmbgNode():
     @classmethod
@@ -16,7 +15,10 @@ class RmbgNode():
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
             "optional": {
-                "content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_input_content_moderation": ("BOOLEAN", {"default": False}), 
+                "visual_output_content_moderation": ("BOOLEAN", {"default": False}), 
+                "preserve_alpha": ("BOOLEAN", {"default": True}), 
+            
             }               
         }
 
@@ -26,10 +28,10 @@ class RmbgNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/background/remove"  # RMBG API URL
+        self.api_url = "https://engine.prod.bria-api.com/v2/image/edit/remove_background"  # RMBG API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, content_moderation, api_key):
+    def execute(self, image, visual_input_content_moderation, visual_output_content_moderation, preserve_alpha, api_key):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -37,31 +39,49 @@ class RmbgNode():
         if isinstance(image, torch.Tensor):
             image = preprocess_image(image)
 
-        # Prepare the API request payload
-        image_buffer = BytesIO()
-        image.save(image_buffer, format="JPEG")
+        # Convert image to base64 for the new API format
+        image_base64 = image_to_base64(image)
+        payload = {
+            "image": image_base64,
+            "visual_input_content_moderation": visual_input_content_moderation,
+            "visual_output_content_moderation":visual_output_content_moderation,
+            "preserve_alpha":preserve_alpha
+        }
 
-        # Get binary data from buffer
-        image_buffer.seek(0)  # Move cursor to the start of the buffer
-        binary_data = image_buffer.read()
-
-        files=[('file',('temp_img.jpeg',  BytesIO(binary_data),'image/jpeg'))]
-        payload = {"content_moderation": content_moderation}   
+        headers = {
+            "Content-Type": "application/json",
+            "api_token": f"{api_key}"
+        }
 
         try:
-            response = requests.post(self.api_url, data=payload, headers={"api_token": api_key}, files=files)
-            # Check for successful response
-            if response.status_code == 200:
-                print('response is 200')
-                # Process the output image from API response
+            response = requests.post(self.api_url, json=payload, headers=headers)
+            
+            if response.status_code == 200 or response.status_code == 202:
+                print('Initial RMBG request successful, polling for completion...')
                 response_dict = response.json()
-                image_response = requests.get(response_dict['result_url'])
+
+                status_url = response_dict.get('status_url')
+                request_id = response_dict.get('request_id')
+                
+                if not status_url:
+                    raise Exception("No status_url returned from API")
+                
+                print(f"Request ID: {request_id}, Status URL: {status_url}")
+                
+                final_response = poll_status_until_completed(status_url, api_key)
+                
+                # Get the result image URL
+                result_image_url = final_response['result']['image_url']
+                
+                # Download and process the result image
+                image_response = requests.get(result_image_url)
                 result_image = Image.open(io.BytesIO(image_response.content))
                 result_image = np.array(result_image).astype(np.float32) / 255.0
                 result_image = torch.from_numpy(result_image)[None,]
+                
                 return (result_image,)
             else:
-                raise Exception(f"Error: API request failed with status code {response.status_code}")
+                raise Exception(f"Error: API request failed with status code {response.status_code} {response.text}")
 
         except Exception as e:
             raise Exception(f"{e}")

nodes/shot_by_text_node.py

@@ -10,7 +10,7 @@ class ShotByTextNode():
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
                 "scene_description": ("STRING",),
-                "optimize_description": ("INT", {"default": 1}),
+                "mode": (["base", "fast", "high_control"], {"default": "high_control"}),
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"})  # API Key input with a default value
             },
             "optional": {
@@ -27,7 +27,7 @@ class ShotByTextNode():
         self.api_url = "https://engine.prod.bria-api.com/v1/product/lifestyle_shot_by_text"  # Eraser API URL
         
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, api_key, scene_description, optimize_description, content_moderation):
+    def execute(self, image, api_key, scene_description, mode, content_moderation):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -35,12 +35,11 @@ class ShotByTextNode():
         if isinstance(image, torch.Tensor):
             image = preprocess_image(image)
 
-        optimize_description = bool(optimize_description)
         image_base64 = image_to_base64(image)
         payload = {
             "file": image_base64,
             "scene_description": scene_description,
-            "optimize_description": optimize_description,
+            "mode": mode,
             "placement_type": "original",
             "original_quality": True,
             "sync": True,

nodes/tailored_portrait_node.py

@@ -0,0 +1,75 @@
+import numpy as np
+import requests
+from PIL import Image
+import io
+import torch
+
+from .common import image_to_base64, preprocess_image
+
+class TailoredPortraitNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "image": ("IMAGE",),  # Input image from another node
+                "tailored_model_id": ("INT",),
+                "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
+            },
+            "optional": {
+                "seed": ("INT", {"default": 123456}),
+                "tailored_model_influence": ("FLOAT", {"default": 0.9}),
+                "id_strength": ("FLOAT", {"default": 0.7}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image",)
+    CATEGORY = "API Nodes"
+    FUNCTION = "execute"  # This is the method that will be executed
+
+    def __init__(self):
+        self.api_url = "https://engine.prod.bria-api.com/v1/tailored-gen/restyle_portrait"  # Eraser API URL
+
+    # Define the execute method as expected by ComfyUI
+    def execute(self, image, tailored_model_id, api_key, seed, tailored_model_influence, id_strength):
+        if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
+            raise Exception("Please insert a valid API key.")
+
+        # Convert the image and mask directly to  if isinstance(image, torch.Tensor):
+        if isinstance(image, torch.Tensor):
+            image = preprocess_image(image)
+            
+        image_base64 = image_to_base64(image)
+
+        # Prepare the API request payload
+        payload = {
+            "id_image_file": f"{image_base64}",
+            "tailored_model_id": tailored_model_id,
+            "tailored_model_influence": tailored_model_influence,
+            "id_strength": id_strength,
+            "seed": seed
+        }
+
+        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['image_res'])
+                result_image = Image.open(io.BytesIO(image_response.content))
+                result_image = result_image.convert("RGB")
+                result_image = np.array(result_image).astype(np.float32) / 255.0
+                result_image = torch.from_numpy(result_image)[None,]
+                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}")

nodes/text_2_image_base_node.py

@@ -38,7 +38,7 @@ class Text2ImageBaseNode():
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        self.api_url = "https://engine.prod.bria-api.com/v1/text-to-image/base/2.3" #"http://0.0.0.0:5000/v1/text-to-image/base/2.3"
+        self.api_url = "https://engine.prod.bria-api.com/v1/text-to-image/base/3.2" 
 
     def execute(
             self, api_key, prompt, aspect_ratio, seed, negative_prompt, 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-bria-api"
 description = "Custom nodes for ComfyUI using BRIA's API."
-version = "2.0.2"
+version = "2.1.0"
 license = {file = "LICENSE"}
 
 [project.urls]