return model id from TG info node

* base and hd

* image prompt

* reimagine

Readme.md

@@ -21,9 +21,42 @@ To load a workflow, import the compatible workflow.json files from this [folder]
 
 # Coming soon
 
-- [ ] Image Generation
 - [ ] Video Editing
 
+# Available Nodes
+
+## Image Generation Nodes
+
+### **Text2Image Base Node**
+This node generates images from text prompts, serving as the foundation for creating visuals based on descriptive input. [[🤗model card](https://huggingface.co/briaai/BRIA-2.3)]
+
+## 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)].
+
+### **Tailored Gen**
+This node generates images using a trained tailored model, faithfully reproducing specific visual IP elements or guidelines established during model training.
+
+### **Tailored Model Info**
+This node retrieves the **default settings** and **prompt prefix** of a **trained tailored model**. It provides the necessary information to configure and run the model in the **Tailored Gen node**, ensuring consistency with the model's intended behavior.  
+
+## 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)]
+
+### **Eraser**
+This node is used to remove specific objects or areas from an image by providing a mask. Powered by BRIA's ControlNet inpainting [[🤗model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Inpainting)] [[🤗HF demo](https://huggingface.co/spaces/briaai/BRIA-Eraser-API)].
+
+### **GenFill**
+This node is used to generate objects by prompt in a specific region of an image. This functionality is powered by BRIA's ControlNet Generative Fill. [[🤗model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Generative-Fill)] [[🤗HF demo](https://huggingface.co/spaces/briaai/BRIA-Generative-Fill-API)]
+
+## 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)]
+
+### **ShotByText**
+This node is used to modify the background in an image by providing a prompt, This functionality is powered by BRIA's ControlNet Background-Generation.[[🤗ContrlNet model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-BG-Gen)] [[🤗HF demo](https://huggingface.co/spaces/briaai/Product-Shot-Generation)].
+
+### **ShotByImage**
+This node is used to modify the background in an image by providing a reference image. This functionality is powered by BRIA's ControlNet Background-Generation and BRIA's Image-Prompt. [[🤗ContrlNet model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Inpainting)] [[🤗IP-Adapter model card](https://huggingface.co/briaai/Image-Prompt)] [[🤗HF demo](https://huggingface.co/spaces/briaai/Product-Shot-Generation)].
+
 # Installation
 There are two methods to install the BRIA ComfyUI API nodes:
 
@@ -44,35 +77,5 @@ There are two methods to install the BRIA ComfyUI API nodes:
 
 3. Restart ComfyUI and load the workflows.
 
-# Available Nodes
-
-## Tailored Generation Nodes
-These nodes use pre-trained tailored models to generate images in a specific visual style based on provided samples. [[API docs](https://bria-ai-api-docs.redoc.ly/tag/Tailored-Generation)].
-
-### **Tailored Model Info**
-This node retrieves the **default settings** and **prompt prefix** of a **trained tailored model**. It provides the necessary information to configure and run the model in the **Tailored Gen node**, ensuring consistency with the model's intended behavior.  
-
-### **Tailored Gen**
-This node is used to generate using a trained tailored model. It is designed to preserve the visual characteristics and ensure style fidelity established during model training.
-
-## 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)]
-
-### **Eraser**
-This node is used to remove specific objects or areas from an image by providing a mask. Powered by BRIA's ControlNet inpainting [[🤗model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Inpainting)] [[🤗HF demo](https://huggingface.co/spaces/briaai/BRIA-Eraser-API)].
-
-### **GenFill**
-This node is used to generate objects by prompt in a specific region of an image. This functionality is powered by BRIA's ControlNet Generative Fill. [[🤗model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Generative-Fill)] [[🤗HF demo](https://huggingface.co/spaces/briaai/BRIA-Generative-Fill-API)]
-
-## 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)]
-
-### **ShotByText**
-This node is used to modify the background in an image by providing a prompt, This functionality is powered by BRIA's ControlNet Background-Generation.[[🤗ContrlNet model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-BG-Gen)] [[🤗HF demo](https://huggingface.co/spaces/briaai/Product-Shot-Generation)].
-
-### **ShotByImage**
-This node is used to modify the background in an image by providing a reference image. This functionality is powered by BRIA's ControlNet Background-Generation and BRIA's Image-Prompt. [[🤗ContrlNet model card](https://huggingface.co/briaai/BRIA-2.3-ControlNet-Inpainting)] [[🤗IP-Adapter model card](https://huggingface.co/briaai/Image-Prompt)] [[🤗HF demo](https://huggingface.co/spaces/briaai/Product-Shot-Generation)].
-
-
 <!-- ### Campaign generation
 Coming soon -->

__init__.py

@@ -1,10 +1,18 @@
-from .nodes import EraserNode, GenFillNode, ShotByTextNode, ShotByImageNode
+from .nodes import (EraserNode, GenFillNode, ShotByTextNode, ShotByImageNode, TailoredGenNode,
+                     TailoredModelInfoNode, Text2ImageBaseNode, Text2ImageFastNode, Text2ImageHDNode,
+                     ReimagineNode)
 # Map the node class to a name used internally by ComfyUI
 NODE_CLASS_MAPPINGS = {
     "BriaEraser": EraserNode,  # Return the class, not an instance
     "BriaGenFill": GenFillNode,
     "ShotByTextNode": ShotByTextNode,
     "ShotByImageNode": ShotByImageNode,
+    "BriaTailoredGen": TailoredGenNode,
+    "TailoredModelInfoNode": TailoredModelInfoNode,
+    "Text2ImageBaseNode": Text2ImageBaseNode,
+    "Text2ImageFastNode": Text2ImageFastNode,
+    "Text2ImageHDNode": Text2ImageHDNode,
+    "ReimagineNode": ReimagineNode,
 }
 # Map the node display name to the one shown in the ComfyUI node interface
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -12,4 +20,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "BriaGenFill": "Bria GenFill",
     "ShotByTextNode": "Bria Shot By Text",
     "ShotByImageNode": "Bria Shot By Image",
+    "BriaTailoredGen": "Bria Tailored Gen",
+    "TailoredModelInfoNode": "Bria Tailored Model Info",
+    "Text2ImageBaseNode": "Bria Text2Image Base",
+    "Text2ImageFastNode": "Bria Text2Image Fast",
+    "Text2ImageHDNode": "Bria Text2Image HD",
+    "ReimagineNode": "Bria Reimagine",
 }

nodes/__init__.py

@@ -2,3 +2,9 @@ from .eraser_node import EraserNode
 from .generative_fill_node import GenFillNode
 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 .text_2_image_base_node import Text2ImageBaseNode
+from .text_2_image_fast_node import Text2ImageFastNode
+from .text_2_image_hd_node import Text2ImageHDNode
+from .reimagine_node import ReimagineNode

nodes/base_node.py

@@ -1,96 +0,0 @@
-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 postprocess_image(self, image):
-        result_image = Image.open(io.BytesIO(image))
-        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
-
-
-    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}")

nodes/common.py

@@ -0,0 +1,92 @@
+import numpy as np
+from PIL import Image
+import io
+import torch
+import base64
+from torchvision.transforms import ToPILImage
+import requests
+
+def postprocess_image(image):
+    result_image = Image.open(io.BytesIO(image))
+    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
+
+def image_to_base64(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 preprocess_image(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(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 process_request(api_url, 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 = preprocess_image(image)
+    if isinstance(mask, torch.Tensor):
+        mask = preprocess_mask(mask)
+
+    # Convert the image and mask directly to Base64 strings
+    image_base64 = image_to_base64(image)
+    mask_base64 = 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(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}")

nodes/eraser_node.py

@@ -1,17 +1,8 @@
-import numpy as np
+from .common import process_request
-import requests
-from PIL import Image
-import io
-import base64
-from torchvision.transforms import ToPILImage, ToTensor
-import torch
 
-from .base_node import BriaAPINode
+class EraserNode():
-
+    @classmethod
-# Eraser Node
+    def INPUT_TYPES(self):
-class EraserNode(BriaAPINode):
-    @staticmethod
-    def INPUT_TYPES():
         return {
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
@@ -26,9 +17,9 @@ class EraserNode(BriaAPINode):
     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
+        self.api_url = "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)
+        return process_request(self.api_url, image, mask, api_key)
     

nodes/generative_fill_node.py

@@ -2,17 +2,14 @@ import numpy as np
 import requests
 from PIL import Image
 import io
-import base64
-from torchvision.transforms import ToPILImage, ToTensor
 import torch
 
-from .base_node import BriaAPINode
+from .common import image_to_base64, preprocess_image, preprocess_mask
 
 
-# Generative Fill Node
+class GenFillNode():
-class GenFillNode(BriaAPINode):
+    @classmethod
-    @staticmethod
+    def INPUT_TYPES(self):
-    def INPUT_TYPES():
         return {
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
@@ -28,7 +25,7 @@ class GenFillNode(BriaAPINode):
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        super().__init__("https://engine.prod.bria-api.com/v1/gen_fill")  # Eraser API URL
+        self.api_url = "https://engine.prod.bria-api.com/v1/gen_fill"  # Eraser API URL
 
     # Define the execute method as expected by ComfyUI
     def execute(self, image, mask, prompt, api_key):
@@ -37,13 +34,13 @@ class GenFillNode(BriaAPINode):
 
         # Check if image and mask are tensors, if so, convert to NumPy arrays
         if isinstance(image, torch.Tensor):
-            image = self.preprocess_image(image)
+            image = preprocess_image(image)
         if isinstance(mask, torch.Tensor):
-            mask = self.preprocess_mask(mask)
+            mask = preprocess_mask(mask)
 
         # Convert the image and mask directly to Base64 strings
-        image_base64 = self.image_to_base64(image)
+        image_base64 = image_to_base64(image)
-        mask_base64 = self.image_to_base64(mask)
+        mask_base64 = image_to_base64(mask)
 
         # Prepare the API request payload
         payload = {

nodes/reimagine_node.py

@@ -0,0 +1,67 @@
+import requests
+
+from .common import postprocess_image, preprocess_image, image_to_base64
+
+
+class ReimagineNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "api_key": ("STRING", ),
+                "prompt": ("STRING",),
+            },
+            "optional": {
+                "seed": ("INT", {"default": -1}),
+                "steps_num": ("INT", {"default": 12}), # if used with tailored, possibly get this from the tailored model info node
+                "structure_ref_influence": ("FLOAT", {"default": 0.75}),
+                "fast": ("INT", {"default": 0}), # if used with tailored, possibly get this from the tailored model info node
+                "structure_image": ("IMAGE", ),
+                "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
+            }
+        }
+
+    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/reimagine" #"http://0.0.0.0:5000/v1/reimagine"
+
+    def execute(
+            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)
+        payload = {
+            "prompt": tailored_generation_prefix + prompt,
+            "num_results": 1,
+            "sync": True,
+            "seed": seed,
+            "steps_num": steps_num,
+            "include_generation_prefix": False,
+        }
+        if structure_image is not None:
+            structure_image = preprocess_image(structure_image)
+            structure_image = image_to_base64(structure_image)
+            payload["structure_image_file"] = structure_image
+            payload["structure_ref_influence"] = structure_ref_influence
+        if tailored_model_id is not None and tailored_model_id != "":
+            payload["tailored_model_id"] = tailored_model_id
+            payload["tailored_model_influence"] = tailored_model_influence
+        response = requests.post(
+            self.api_url,
+            json=payload,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result'][0]["urls"][0])
+                result_image = postprocess_image(image_response.content)
+                return (result_image,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

nodes/shot_by_image_node.py

@@ -1,17 +1,11 @@
-import numpy as np
 import requests
-from PIL import Image
-import io
-import base64
-from torchvision.transforms import ToPILImage, ToTensor
 import torch
 
-from .base_node import BriaAPINode
+from .common import postprocess_image, preprocess_image, image_to_base64
 
-# shot by image Node
+class ShotByImageNode():
-class ShotByImageNode(BriaAPINode):
+    @classmethod
-    @staticmethod
+    def INPUT_TYPES(self):
-    def INPUT_TYPES():
         return {
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
@@ -27,7 +21,7 @@ class ShotByImageNode(BriaAPINode):
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        super().__init__("https://engine.prod.bria-api.com/v1/product/lifestyle_shot_by_image")  # Eraser API URL
+        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, ):
@@ -36,13 +30,13 @@ class ShotByImageNode(BriaAPINode):
 
         # Check if image and mask are tensors, if so, convert to NumPy arrays
         if isinstance(image, torch.Tensor):
-            image = self.preprocess_image(image)
+            image = preprocess_image(image)
         if isinstance(ref_image, torch.Tensor):
-            ref_image = self.preprocess_image(ref_image)
+            ref_image = preprocess_image(ref_image)
 
         # Convert the image and mask directly to Base64 strings                    
-        image_base64 = self.image_to_base64(image)
+        image_base64 = image_to_base64(image)
-        ref_image_base64 = self.image_to_base64(ref_image)
+        ref_image_base64 = image_to_base64(ref_image)
         enhance_ref_image = bool(enhance_ref_image)
 
         payload = {
@@ -65,7 +59,7 @@ class ShotByImageNode(BriaAPINode):
                 # Process the output image from API response
                 response_dict = response.json()
                 image_response = requests.get(response_dict['result'][0][0])
-                result_image = self.postprocess_image(image_response.content)
+                result_image = postprocess_image(image_response.content)
                 return (result_image,)
             else:
                 raise Exception(f"Error: API request failed with status code {response.status_code}")

nodes/shot_by_text_node.py

@@ -1,17 +1,11 @@
-import numpy as np
 import requests
-from PIL import Image
-import io
-import base64
-from torchvision.transforms import ToPILImage, ToTensor
 import torch
 
-from .base_node import BriaAPINode
+from .common import postprocess_image, preprocess_image, image_to_base64
 
-# shot by text Node
+class ShotByTextNode():
-class ShotByTextNode(BriaAPINode):
+    @classmethod
-    @staticmethod
+    def INPUT_TYPES(self):
-    def INPUT_TYPES():
         return {
             "required": {
                 "image": ("IMAGE",),  # Input image from another node
@@ -27,8 +21,8 @@ class ShotByTextNode(BriaAPINode):
     FUNCTION = "execute"  # This is the method that will be executed
 
     def __init__(self):
-        super().__init__("https://engine.prod.bria-api.com/v1/product/lifestyle_shot_by_text")  # Eraser API URL
+        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, ):
         if api_key.strip() == "" or api_key.strip() == "BRIA_API_TOKEN":
@@ -36,10 +30,10 @@ class ShotByTextNode(BriaAPINode):
 
         # Check if image and mask are tensors, if so, convert to NumPy arrays
         if isinstance(image, torch.Tensor):
-            image = self.preprocess_image(image)
+            image = preprocess_image(image)
 
         optimize_description = bool(optimize_description)
-        image_base64 = self.image_to_base64(image)
+        image_base64 = image_to_base64(image)
         payload = {
             "file": image_base64,
             "scene_description": scene_description,
@@ -60,7 +54,7 @@ class ShotByTextNode(BriaAPINode):
                 # Process the output image from API response
                 response_dict = response.json()
                 image_response = requests.get(response_dict['result'][0][0])
-                result_image = self.postprocess_image(image_response.content)
+                result_image = postprocess_image(image_response.content)
                 return (result_image,)
             else:
                 raise Exception(f"Error: API request failed with status code {response.status_code}")

nodes/tailored_gen_node.py

@@ -0,0 +1,82 @@
+import requests
+
+from .common import postprocess_image, preprocess_image, image_to_base64
+
+
+class TailoredGenNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "model_id": ("STRING",),
+                "api_key": ("STRING", ),
+            },
+            "optional": {
+                "prompt": ("STRING",),
+                "generation_prefix": ("STRING",), # possibly get this from the tailored model info node
+                "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}),
+                "seed": ("INT", {"default": -1}),
+                "model_influence": ("FLOAT", {"default": 1.0}),
+                "negative_prompt": ("STRING", {"default": ""}),
+                "fast": ("INT", {"default": 1}), # possibly get this from the tailored model info node
+                "steps_num": ("INT", {"default": 8}), # possibly get this from the tailored model info node
+                "guidance_method_1": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}),
+                "guidance_method_1_scale": ("FLOAT", {"default": 1.0}),
+                "guidance_method_1_image": ("IMAGE", ),
+                "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", ),
+            }
+        }
+
+    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/text-to-image/tailored/" #"http://0.0.0.0:5000/v1/text-to-image/tailored/"
+
+    def execute(
+            self, model_id, api_key, prompt, generation_prefix, aspect_ratio, 
+            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,
+        ):
+        fast = bool(fast)
+        payload = {
+            "prompt": generation_prefix + prompt,
+            "num_results": 1,
+            "aspect_ratio": aspect_ratio,
+            "sync": True,
+            "seed": seed,
+            "model_influence": model_influence,
+            "negative_prompt": negative_prompt,
+            "fast": fast,
+            "steps_num": steps_num,
+            "include_generation_prefix": False,
+        }
+        if guidance_method_1_image is not None:
+            guidance_method_1_image = preprocess_image(guidance_method_1_image)
+            guidance_method_1_image = image_to_base64(guidance_method_1_image)
+            payload["guidance_method_1"] = guidance_method_1
+            payload["guidance_method_1_scale"] = guidance_method_1_scale
+            payload["guidance_method_1_image_file"] = guidance_method_1_image
+        if guidance_method_2_image is not None:
+            guidance_method_2_image = preprocess_image(guidance_method_2_image)
+            guidance_method_2_image = image_to_base64(guidance_method_2_image)
+            payload["guidance_method_2"] = guidance_method_2
+            payload["guidance_method_2_scale"] = guidance_method_2_scale
+            payload["guidance_method_2_image_file"] = guidance_method_2_image
+        response = requests.post(
+            self.api_url + model_id,
+            json=payload,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result'][0]["urls"][0])
+                result_image = postprocess_image(image_response.content)
+                return (result_image,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

nodes/tailored_model_info_node.py

@@ -0,0 +1,35 @@
+import requests
+
+
+class TailoredModelInfoNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "model_id": ("STRING",),
+                "api_key": ("STRING", )
+            }
+        }
+
+    RETURN_TYPES = ("STRING", "STRING","INT", "INT", )
+    RETURN_NAMES = ("generation_prefix", "model_id", "default_fast",  "default_steps_num", )
+    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/models/"
+
+    # Define the execute method as expected by ComfyUI
+    def execute(self, model_id, api_key):
+        response = requests.get(
+            self.api_url + model_id,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+            generation_prefix = response.json()["generation_prefix"]
+            training_version = response.json()["training_version"]
+            default_fast = 1 if training_version == "light" else 0
+            default_steps_num = 8 if training_version == "light" else 30
+            return (generation_prefix, model_id, default_fast, default_steps_num,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

nodes/text_2_image_base_node.py

@@ -0,0 +1,92 @@
+import requests
+
+from .common import postprocess_image, preprocess_image, image_to_base64
+
+
+class Text2ImageBaseNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "api_key": ("STRING", ),
+            },
+            "optional": {
+                "prompt": ("STRING",),
+                "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}),
+                "seed": ("INT", {"default": -1}),
+                "negative_prompt": ("STRING", {"default": ""}),
+                "steps_num": ("INT", {"default": 30}), 
+                "prompt_enhancement": ("INT", {"default": 0}),
+                "text_guidance_scale": ("INT", {"default": 5}),
+                "medium": (["photography", "art", "none"], {"default": "none"}),
+                "guidance_method_1": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}),
+                "guidance_method_1_scale": ("FLOAT", {"default": 1.0}),
+                "guidance_method_1_image": ("IMAGE", ),
+                "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", ),
+                "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}),
+                "image_prompt_image": ("IMAGE", ),
+                "image_prompt_scale": ("FLOAT", {"default": 1.0}),
+            }
+        }
+
+    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/text-to-image/base/2.3" #"http://0.0.0.0:5000/v1/text-to-image/base/2.3"
+
+    def execute(
+            self, api_key, prompt, aspect_ratio, seed, negative_prompt, 
+            steps_num, prompt_enhancement, text_guidance_scale, medium,
+            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,
+        ):
+        prompt_enhancement = bool(prompt_enhancement)
+        payload = {
+            "prompt": prompt,
+            "num_results": 1,
+            "aspect_ratio": aspect_ratio,
+            "sync": True,
+            "seed": seed,
+            "negative_prompt": negative_prompt,
+            "steps_num": steps_num,
+            "text_guidance_scale": text_guidance_scale,
+            "prompt_enhancement": prompt_enhancement,
+        }
+        if medium != "none":
+            payload["medium"] = medium
+        if guidance_method_1_image is not None:
+            guidance_method_1_image = preprocess_image(guidance_method_1_image)
+            guidance_method_1_image = image_to_base64(guidance_method_1_image)
+            payload["guidance_method_1"] = guidance_method_1
+            payload["guidance_method_1_scale"] = guidance_method_1_scale
+            payload["guidance_method_1_image_file"] = guidance_method_1_image
+        if guidance_method_2_image is not None:
+            guidance_method_2_image = preprocess_image(guidance_method_2_image)
+            guidance_method_2_image = image_to_base64(guidance_method_2_image)
+            payload["guidance_method_2"] = guidance_method_2
+            payload["guidance_method_2_scale"] = guidance_method_2_scale
+            payload["guidance_method_2_image_file"] = guidance_method_2_image
+        if image_prompt_image is not None:
+            image_prompt_image = preprocess_image(image_prompt_image)
+            image_prompt_image = image_to_base64(image_prompt_image)
+            payload["image_prompt_mode"] = image_prompt_mode
+            payload["image_prompt_file"] = image_prompt_image
+            payload["image_prompt_scale"] = image_prompt_scale
+        response = requests.post(
+            self.api_url,
+            json=payload,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result'][0]["urls"][0])
+                result_image = postprocess_image(image_response.content)
+                return (result_image,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

nodes/text_2_image_fast_node.py

@@ -0,0 +1,85 @@
+import requests
+
+from .common import postprocess_image, preprocess_image, image_to_base64
+
+
+class Text2ImageFastNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "api_key": ("STRING", ),
+            },
+            "optional": {
+                "prompt": ("STRING",),
+                "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}),
+                "seed": ("INT", {"default": -1}),
+                "steps_num": ("INT", {"default": 8}), 
+                "prompt_enhancement": ("INT", {"default": 0}),
+                "guidance_method_1": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}),
+                "guidance_method_1_scale": ("FLOAT", {"default": 1.0}),
+                "guidance_method_1_image": ("IMAGE", ),
+                "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", ),
+                "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}),
+                "image_prompt_image": ("IMAGE", ),
+                "image_prompt_scale": ("FLOAT", {"default": 1.0}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image",)
+    CATEGORY = "API Nodes"
+    FUNCTION = "execute"
+
+    def __init__(self):
+        self.api_url = "https://engine.prod.bria-api.com/v1/text-to-image/fast/2.3" #"http://0.0.0.0:5000/v1/text-to-image/fast/2.3"
+
+    def execute(
+            self, api_key, prompt, aspect_ratio, seed, 
+            steps_num, prompt_enhancement,
+            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,
+        ):
+        prompt_enhancement = bool(prompt_enhancement)
+        payload = {
+            "prompt": prompt,
+            "num_results": 1,
+            "aspect_ratio": aspect_ratio,
+            "sync": True,
+            "seed": seed,
+            "steps_num": steps_num,
+            "prompt_enhancement": prompt_enhancement,
+        }
+        if guidance_method_1_image is not None:
+            guidance_method_1_image = preprocess_image(guidance_method_1_image)
+            guidance_method_1_image = image_to_base64(guidance_method_1_image)
+            payload["guidance_method_1"] = guidance_method_1
+            payload["guidance_method_1_scale"] = guidance_method_1_scale
+            payload["guidance_method_1_image_file"] = guidance_method_1_image
+        if guidance_method_2_image is not None:
+            guidance_method_2_image = preprocess_image(guidance_method_2_image)
+            guidance_method_2_image = image_to_base64(guidance_method_2_image)
+            payload["guidance_method_2"] = guidance_method_2
+            payload["guidance_method_2_scale"] = guidance_method_2_scale
+            payload["guidance_method_2_image_file"] = guidance_method_2_image
+        if image_prompt_image is not None:
+            image_prompt_image = preprocess_image(image_prompt_image)
+            image_prompt_image = image_to_base64(image_prompt_image)
+            payload["image_prompt_mode"] = image_prompt_mode
+            payload["image_prompt_file"] = image_prompt_image
+            payload["image_prompt_scale"] = image_prompt_scale
+        response = requests.post(
+            self.api_url,
+            json=payload,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result'][0]["urls"][0])
+                result_image = postprocess_image(image_response.content)
+                return (result_image,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

nodes/text_2_image_hd_node.py

@@ -0,0 +1,62 @@
+import requests
+
+from .common import postprocess_image
+
+
+class Text2ImageHDNode():
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "api_key": ("STRING", ),
+            },
+            "optional": {
+                "prompt": ("STRING",),
+                "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}),
+                "seed": ("INT", {"default": -1}),
+                "negative_prompt": ("STRING", {"default": ""}),
+                "steps_num": ("INT", {"default": 30}), 
+                "prompt_enhancement": ("INT", {"default": 0}),
+                "text_guidance_scale": ("INT", {"default": 5}),
+                "medium": (["photography", "art", "none"], {"default": "none"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image",)
+    CATEGORY = "API Nodes"
+    FUNCTION = "execute"
+
+    def __init__(self):
+        self.api_url = "https://engine.prod.bria-api.com/v1/text-to-image/hd/2.3" #"http://0.0.0.0:5000/v1/text-to-image/hd/2.3"
+
+    def execute(
+            self, api_key, prompt, aspect_ratio, seed, negative_prompt, 
+            steps_num, prompt_enhancement, text_guidance_scale, medium,
+        ):
+        prompt_enhancement = bool(prompt_enhancement)
+        payload = {
+            "prompt": prompt,
+            "num_results": 1,
+            "aspect_ratio": aspect_ratio,
+            "sync": True,
+            "seed": seed,
+            "negative_prompt": negative_prompt,
+            "steps_num": steps_num,
+            "text_guidance_scale": text_guidance_scale,
+            "prompt_enhancement": prompt_enhancement,
+        }
+        if medium != "none":
+            payload["medium"] = medium
+        response = requests.post(
+            self.api_url,
+            json=payload,
+            headers={"api_token": api_key}
+        )
+        if response.status_code == 200:
+                response_dict = response.json()
+                image_response = requests.get(response_dict['result'][0]["urls"][0])
+                result_image = postprocess_image(image_response.content)
+                return (result_image,)
+        else:
+            raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}")

pyproject.toml

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

workflows/tailored_workflow.json

@@ -0,0 +1 @@
+{"last_node_id":21,"last_link_id":49,"nodes":[{"id":2,"type":"TailoredModelInfoNode","pos":[480.2208557128906,641.4290161132812],"size":[315,122],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"generation_prefix","type":"STRING","links":[2],"slot_index":0,"localized_name":"generation_prefix"},{"name":"default_fast","type":"INT","links":[46],"slot_index":1,"localized_name":"default_fast"},{"name":"default_steps_num","type":"INT","links":[40],"slot_index":2,"localized_name":"default_steps_num"}],"properties":{"Node name for S&R":"TailoredModelInfoNode"},"widgets_values":["",""]},{"id":3,"type":"PreviewImage","pos":[1728.9140625,688.2759399414062],"size":[210,246],"flags":{},"order":6,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":41,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"LoadImage","pos":[693.87060546875,831.6130981445312],"size":[315,314],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[48],"slot_index":0,"localized_name":"IMAGE"},{"name":"MASK","type":"MASK","links":null,"localized_name":"MASK"}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":5,"type":"JjkShowText","pos":[847.2867431640625,591.890625],"size":[315,76],"flags":{},"order":4,"mode":0,"inputs":[{"name":"text","type":"STRING","link":2,"widget":{"name":"text"}}],"outputs":[{"name":"text","type":"STRING","links":[49],"slot_index":0,"shape":6,"localized_name":"text"}],"properties":{"Node name for S&R":"JjkShowText"},"widgets_values":["A photo of a character named Sami, a siamese cat with blue eyes, "]},{"id":15,"type":"BriaTailoredGen","pos":[1207.595947265625,645.6781005859375],"size":[456,438],"flags":{},"order":5,"mode":0,"inputs":[{"name":"guidance_method_1_image","type":"IMAGE","link":48,"shape":7,"localized_name":"guidance_method_1_image"},{"name":"guidance_method_2_image","type":"IMAGE","link":null,"shape":7,"localized_name":"guidance_method_2_image"},{"name":"generation_prefix","type":"STRING","link":49,"widget":{"name":"generation_prefix"},"shape":7},{"name":"fast","type":"INT","link":46,"widget":{"name":"fast"},"shape":7},{"name":"steps_num","type":"INT","link":40,"widget":{"name":"steps_num"},"shape":7}],"outputs":[{"name":"output_image","type":"IMAGE","links":[41],"slot_index":0,"localized_name":"output_image"}],"properties":{"Node name for S&R":"BriaTailoredGen"},"widgets_values":["","","a cat","","4:3",-1,"randomize",1,"","",1,"controlnet_canny",1,"controlnet_canny",1]},{"id":21,"type":"Note","pos":[1215.2469482421875,522.4407348632812],"size":[449.75360107421875,58],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["You can get your BRIA API token at: https://bria.ai/api/"],"color":"#432","bgcolor":"#653"},{"id":19,"type":"Note","pos":[484.5993957519531,486.4328918457031],"size":[306.0655212402344,89.87609100341797],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["This node is used to retrieve default settings and prompt prefixes for the chosen tailored model."],"color":"#432","bgcolor":"#653"}],"links":[[2,2,0,5,0,"STRING"],[40,2,2,15,4,"INT"],[41,15,0,3,0,"IMAGE"],[46,2,1,15,3,"INT"],[48,11,0,15,0,"IMAGE"],[49,5,0,15,2,"STRING"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.7627768444385483,"offset":[-122.90473166350671,-300.2018923615813]},"node_versions":{"comfyui-bria-api":"c72754d15b53a13ee0c0419d70401232c56b7fdb","comfy-core":"v0.3.8-1-gc441048","ComfyUI-Jjk-Nodes":"b3c99bb78a99551776b5eab1a820e1cd58f84f31"}},"version":0.4}