From 28c0cdf8d5511fff144e806f66621ec57b49a90d Mon Sep 17 00:00:00 2001
From: israelweiss90 <israel@bria.ai>
Date: Sun, 2 Feb 2025 16:14:02 +0000
Subject: [PATCH] image expansion & remove fg- content moderation

---
 nodes/image_expansion_node.py   | 6 ++++--
 nodes/remove_foreground_node.py | 7 +++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/nodes/image_expansion_node.py b/nodes/image_expansion_node.py
index ef5c86f..4af3e6b 100644
--- a/nodes/image_expansion_node.py
+++ b/nodes/image_expansion_node.py
@@ -23,6 +23,7 @@ class ImageExpansionNode():
                 "prompt": ("STRING", {"default": ""}),
                 "seed": ("INT", {"default": 681794}),
                 "negative_prompt": ("STRING", {"default": "Ugly, mutated"}),
+                "content_moderation": ("BOOLEAN", {"default": False}), 
             } 
         }
 
@@ -42,6 +43,7 @@ class ImageExpansionNode():
                 prompt,
                 seed,
                 negative_prompt,
+                content_moderation,
                 api_key):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
             raise Exception("Please insert a valid API key.")
@@ -70,8 +72,8 @@ class ImageExpansionNode():
             "canvas_size": canvas_size,
             "prompt": prompt,
             "negative_prompt": negative_prompt,
-            "seed": seed
-            # "sync": True
+            "seed": seed,
+            "content_moderation": content_moderation
         }
 
         headers = {
diff --git a/nodes/remove_foreground_node.py b/nodes/remove_foreground_node.py
index b941041..24d8f77 100644
--- a/nodes/remove_foreground_node.py
+++ b/nodes/remove_foreground_node.py
@@ -15,6 +15,9 @@ class RemoveForegroundNode():
                 "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 RemoveForegroundNode():
         self.api_url = "https://engine.internal.prod.bria-api.com/v1/erase_foreground"  # remove foreground 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.")
 
@@ -41,7 +44,7 @@ class RemoveForegroundNode():
         
 #         files=[('file',('temp_img.jpeg', open(temp_img_path, 'rb'),'image/jpeg'))
 #   ]
-        payload = {"file": image_to_base64(image)}
+        payload = {"file": image_to_base64(image), "content_moderation": content_moderation}
 
         headers = {
             "Content-Type": "application/json",