fix image preprocess

content moderation t2i reimagime tailored

Readme.md

@@ -4,14 +4,22 @@
   <img src="./images/Bria Logo.svg" alt="BRIA Logo" width="200"/>
 </p>
 
-This repository provides custom nodes for ComfyUI, enabling direct access to **BRIA's API endpoints** for image generation workflows. **API documentation** is available [**here**](https://bria-ai-api-docs.redoc.ly/#operation//generation/bria-v2/text-to-image). 
+This repository provides custom nodes for ComfyUI, enabling direct access to **BRIA's API endpoints** for image generation and editing workflows. **API documentation** is available [**here**](https://docs.bria.ai/). 
+
+BRIA's APIs and models are built for commercial use and trained on 100% licensed data and does not contain copyrighted materials, such as fictional characters, logos, trademarks, public figures, harmful content, or privacy-infringing content.
 
 An API token is required to use the nodes in your workflows. Get started quickly here
 <a href="https://bria.ai/api/" style="text-decoration:none; vertical-align:middle;">
   <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 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).  
+<p align="center">
+  <img src="./images/background_workflow.png" width="1200"/>
+</p>
 
 
 <!-- Placeholder image of cool workflows. -->
@@ -22,16 +30,17 @@ To load a workflow, import the compatible workflow.json files from this [folder]
 # Available Nodes
 
 ## Image Generation Nodes
-These nodes create high-quality images from text or image prompts, generating photorealistic or artistic results with support for various aspect ratios. [[API docs](https://bria-ai-api-docs.redoc.ly/tag/Image-Generation)].
+These nodes create high-quality images from text or image prompts, generating photorealistic or artistic results with support for various aspect ratios.
 
 | Node                   | Description                                                        |
 |------------------------|--------------------------------------------------------------------|
 | **Text2Image Base**    | Generates images from text prompts, serving as the foundation for text-based image creation. |
 | **Text2Image Fast**    | Optimized for speed, this node generates images from text prompts with faster results while maintaining quality. |
 | **Text2Image HD**      | Optimized for high-resolution outputs, this node generates detailed and sharp visuals from text prompts. |
+| **Reimagine**          | Guides image generation using both prompts and an input image. Preserve the original structure and depth while introducing new materials, colors, and textures. |
 
 ## Tailored Generation Nodes
-These nodes use pre-trained tailored models to generate images that faithfully reproduce specific visual IP elements or guidelines. [[API docs](https://bria-ai-api-docs.redoc.ly/tag/Tailored-Generation)].
+These nodes use pre-trained tailored models to generate images that faithfully reproduce specific visual IP elements or guidelines.
 
 | Node                   | Description                                                        |
 |------------------------|--------------------------------------------------------------------|
@@ -39,7 +48,7 @@ These nodes use pre-trained tailored models to generate images that faithfully r
 | **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. |
 
 ## Image Editing Nodes
-These nodes modify specific parts of images, enabling adjustments while maintaining the integrity of the rest of the image. [[API docs](https://bria-ai-api-docs.redoc.ly/tag/Image-Editing)].
+These nodes modify specific parts of images, enabling adjustments while maintaining the integrity of the rest of the image.
 
 | Node                   | Description                                                        |
 |------------------------|--------------------------------------------------------------------|
@@ -50,8 +59,8 @@ These nodes modify specific parts of images, enabling adjustments while maintain
 | **GenFill**            | Generates objects by prompt in a specific region of an image. |
 | **Erase Foreground**    | Removes the foreground from an image, isolating the background. |
 
-## Product Shot Generation Nodes
-These nodes create high-quality product images for eCommerce workflows. [[API docs](https://bria-ai-api-docs.redoc.ly/tag/Product-Shots-Generation)].
+## Product Shot Editing Nodes
+These nodes create high-quality product images for eCommerce workflows.
 
 | Node                   | Description                                                        |
 |------------------------|--------------------------------------------------------------------|

__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,

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/reimagine_node.py

@@ -20,6 +20,7 @@ class ReimagineNode():
                 "tailored_model_id": ("STRING", ),
                 "tailored_model_influence": ("FLOAT", {"default": 0.5}),
                 "tailored_generation_prefix": ("STRING",), # if used with tailored, possibly get this from the tailored model info node
+                "content_moderation": ("INT", {"default": 0}),
             }
         }
 
@@ -35,8 +36,8 @@ class ReimagineNode():
             self, api_key, prompt, seed, 
             steps_num, fast, structure_ref_influence, structure_image=None,
             tailored_model_id=None, tailored_model_influence=None, tailored_generation_prefix=None,
-        ):
-        fast = bool(fast)
+            content_moderation=0,
+        ):        
         payload = {
             "prompt": tailored_generation_prefix + prompt,
             "num_results": 1,
@@ -44,6 +45,7 @@ class ReimagineNode():
             "seed": seed,
             "steps_num": steps_num,
             "include_generation_prefix": False,
+            "content_moderation": content_moderation,
         }
         if structure_image is not None:
             structure_image = preprocess_image(structure_image)

nodes/replace_bg_node.py

@@ -24,7 +24,8 @@ class ReplaceBgNode():
                 "original_quality": ("BOOLEAN", {"default": False}), 
                 "force_rmbg": ("BOOLEAN", {"default": False}), 
                 "negative_prompt": ("STRING", {"default": None}), 
-                "seed": ("INT", {"default": 681794})
+                "seed": ("INT", {"default": 681794}),
+                "content_moderation": ("BOOLEAN", {"default": False}), 
             }
         }
 
@@ -45,6 +46,7 @@ class ReplaceBgNode():
                 negative_prompt,
                 seed,
                 api_key,
+                content_moderation,
                 bg_prompt=None,
                 ref_image=None,):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
@@ -74,7 +76,8 @@ class ReplaceBgNode():
             "negative_prompt": negative_prompt,
             "seed": seed,
             "sync": True,
-            "num_results": 1
+            "num_results": 1,
+            "content_moderation": content_moderation            
         }   
 
         headers = {

nodes/rmbg_node.py

@@ -15,6 +15,9 @@ class RmbgNode():
                 "image": ("IMAGE",),  # Input image from another node
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"}), # API Key input with a default value
             },
+            "optional": {
+                "content_moderation": ("BOOLEAN", {"default": False}), 
+            }               
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -26,7 +29,7 @@ class RmbgNode():
         self.api_url = "https://engine.prod.bria-api.com/v1/background/remove"  # RMBG API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, api_key):
+    def execute(self, image, content_moderation, api_key):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -43,9 +46,10 @@ class RmbgNode():
         binary_data = image_buffer.read()
 
         files=[('file',('temp_img.jpeg',  BytesIO(binary_data),'image/jpeg'))]
+        payload = {"content_moderation": content_moderation}   
 
         try:
-            response = requests.post(self.api_url, data={}, headers={"api_token": api_key}, files=files)
+            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')

nodes/shot_by_image_node.py

@@ -12,7 +12,10 @@ class ShotByImageNode():
                 "ref_image": ("IMAGE",),  # ref image from another node
                 "enhance_ref_image": ("INT", {"default": 1}),
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"})  # API Key input with a default value
-            }
+            },
+            "optional": {
+                "content_moderation": ("BOOLEAN", {"default": False}), 
+            }        
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -24,7 +27,7 @@ class ShotByImageNode():
         self.api_url = "https://engine.prod.bria-api.com/v1/product/lifestyle_shot_by_image"  # Eraser API URL
 
     # Define the execute method as expected by ComfyUI
-    def execute(self, image, ref_image, api_key, enhance_ref_image, ):
+    def execute(self, image, ref_image, api_key, enhance_ref_image, content_moderation):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -45,7 +48,8 @@ class ShotByImageNode():
             "enhance_ref_image": enhance_ref_image,
             "placement_type": "original",
             "original_quality": True, 
-            "sync": True
+            "sync": True,
+            "content_moderation": content_moderation
         }
         headers = {
             "Content-Type": "application/json",

nodes/shot_by_text_node.py

@@ -12,7 +12,10 @@ class ShotByTextNode():
                 "scene_description": ("STRING",),
                 "optimize_description": ("INT", {"default": 1}),
                 "api_key": ("STRING", {"default": "BRIA_API_TOKEN"})  # API Key input with a default value
-            }
+            },
+            "optional": {
+                "content_moderation": ("BOOLEAN", {"default": False}), 
+            }   
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -24,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, ):
+    def execute(self, image, api_key, scene_description, optimize_description, content_moderation):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
 
@@ -40,7 +43,9 @@ class ShotByTextNode():
             "optimize_description": optimize_description,
             "placement_type": "original",
             "original_quality": True,
-            "sync": True
+            "sync": True,
+            "content_moderation": content_moderation
+
         }
         headers = {
             "Content-Type": "application/json",

nodes/tailored_gen_node.py

@@ -26,6 +26,7 @@ class TailoredGenNode():
                 "guidance_method_2": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}),
                 "guidance_method_2_scale": ("FLOAT", {"default": 1.0}),
                 "guidance_method_2_image": ("IMAGE", ),
+                "content_moderation": ("INT", {"default": 0}),
             }
         }
 
@@ -42,8 +43,8 @@ class TailoredGenNode():
             seed, model_influence, negative_prompt, fast, steps_num,
             guidance_method_1=None, guidance_method_1_scale=None, guidance_method_1_image=None,
             guidance_method_2=None, guidance_method_2_scale=None, guidance_method_2_image=None,
+            content_moderation=0,
         ):
-        fast = bool(fast)
         payload = {
             "prompt": generation_prefix + prompt,
             "num_results": 1,
@@ -55,6 +56,7 @@ class TailoredGenNode():
             "fast": fast,
             "steps_num": steps_num,
             "include_generation_prefix": False,
+            "content_moderation": content_moderation,
         }
         if guidance_method_1_image is not None:
             guidance_method_1_image = preprocess_image(guidance_method_1_image)

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

@@ -28,6 +28,7 @@ class Text2ImageBaseNode():
                 "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}),
                 "image_prompt_image": ("IMAGE", ),
                 "image_prompt_scale": ("FLOAT", {"default": 1.0}),
+                "content_moderation": ("INT", {"default": 0}),
             }
         }
 
@@ -45,8 +46,8 @@ class Text2ImageBaseNode():
             guidance_method_1=None, guidance_method_1_scale=None, guidance_method_1_image=None,
             guidance_method_2=None, guidance_method_2_scale=None, guidance_method_2_image=None,
             image_prompt_mode=None, image_prompt_image=None, image_prompt_scale=None,
+            content_moderation=0,
         ):
-        prompt_enhancement = bool(prompt_enhancement)
         payload = {
             "prompt": prompt,
             "num_results": 1,
@@ -57,6 +58,7 @@ class Text2ImageBaseNode():
             "steps_num": steps_num,
             "text_guidance_scale": text_guidance_scale,
             "prompt_enhancement": prompt_enhancement,
+            "content_moderation": content_moderation,
         }
         if medium != "none":
             payload["medium"] = medium

nodes/text_2_image_fast_node.py

@@ -25,6 +25,7 @@ class Text2ImageFastNode():
                 "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}),
                 "image_prompt_image": ("IMAGE", ),
                 "image_prompt_scale": ("FLOAT", {"default": 1.0}),
+                "content_moderation": ("INT", {"default": 0}),
             }
         }
 
@@ -42,8 +43,8 @@ class Text2ImageFastNode():
             guidance_method_1=None, guidance_method_1_scale=None, guidance_method_1_image=None,
             guidance_method_2=None, guidance_method_2_scale=None, guidance_method_2_image=None,
             image_prompt_mode=None, image_prompt_image=None, image_prompt_scale=None,
+            content_moderation=0,
         ):
-        prompt_enhancement = bool(prompt_enhancement)
         payload = {
             "prompt": prompt,
             "num_results": 1,
@@ -52,6 +53,7 @@ class Text2ImageFastNode():
             "seed": seed,
             "steps_num": steps_num,
             "prompt_enhancement": prompt_enhancement,
+            "content_moderation": content_moderation,
         }
         if guidance_method_1_image is not None:
             guidance_method_1_image = preprocess_image(guidance_method_1_image)

nodes/text_2_image_hd_node.py

@@ -19,6 +19,7 @@ class Text2ImageHDNode():
                 "prompt_enhancement": ("INT", {"default": 0}),
                 "text_guidance_scale": ("INT", {"default": 5}),
                 "medium": (["photography", "art", "none"], {"default": "none"}),
+                "content_moderation": ("INT", {"default": 0}),
             }
         }
 
@@ -32,9 +33,8 @@ class Text2ImageHDNode():
 
     def execute(
             self, api_key, prompt, aspect_ratio, seed, negative_prompt, 
-            steps_num, prompt_enhancement, text_guidance_scale, medium,
+            steps_num, prompt_enhancement, text_guidance_scale, medium, content_moderation=0,
         ):
-        prompt_enhancement = bool(prompt_enhancement)
         payload = {
             "prompt": prompt,
             "num_results": 1,
@@ -45,6 +45,7 @@ class Text2ImageHDNode():
             "steps_num": steps_num,
             "text_guidance_scale": text_guidance_scale,
             "prompt_enhancement": prompt_enhancement,
+            "content_moderation": content_moderation,
         }
         if medium != "none":
             payload["medium"] = medium

pyproject.toml

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

workflows/text_to_image_workflow.json

@@ -0,0 +1,169 @@
+{
+  "last_node_id": 13,
+  "last_link_id": 11,
+  "nodes": [
+    {
+      "id": 12,
+      "type": "LoadImage",
+      "pos": [
+        669.7035522460938,
+        136.97129821777344
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            10
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "pexels-photo-3246665.png",
+        "image"
+      ]
+    },
+    {
+      "id": 11,
+      "type": "PreviewImage",
+      "pos": [
+        1587.27001953125,
+        97.39167022705078
+      ],
+      "size": [
+        210,
+        246
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 11
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    },
+    {
+      "id": 10,
+      "type": "Text2ImageFastNode",
+      "pos": [
+        1058.320556640625,
+        96.64427185058594
+      ],
+      "size": [
+        438.71258544921875,
+        394.27716064453125
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "guidance_method_1_image",
+          "type": "IMAGE",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "guidance_method_2_image",
+          "type": "IMAGE",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "image_prompt_image",
+          "type": "IMAGE",
+          "link": 10,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "output_image",
+          "type": "IMAGE",
+          "links": [
+            11
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Text2ImageFastNode"
+      },
+      "widgets_values": [
+        "BRIA_API_TOKEN",
+        "A drawing of a lion on a table.\t",
+        "4:3",
+        990,
+        "randomize",
+        8,
+        0,
+        "controlnet_canny",
+        1,
+        "controlnet_canny",
+        1,
+        "regular",
+        1
+      ]
+    }
+  ],
+  "links": [
+    [
+      10,
+      12,
+      0,
+      10,
+      2,
+      "IMAGE"
+    ],
+    [
+      11,
+      10,
+      0,
+      11,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.7513148009015777,
+      "offset": [
+        11.112206386364164,
+        66.47311795454547
+      ]
+    },
+    "node_versions": {
+      "comfy-core": "v0.3.10-42-gff83865",
+      "comfyui-bria-api": "499ec5d104cc5110407eafce468ce1d47ac168b3"
+    },
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0
+  },
+  "version": 0.4
+}