From a970f19bf2398cdff2dcb6a3a31072b5bd275a0d Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 13 Jan 2026 19:38:15 +0000 Subject: [PATCH] WAI-4383 --- __init__.py | 12 +- nodes/__init__.py | 3 + nodes/fibo_edit_node.py | 163 ++++++++++++++++++ .../fibo_edit_structured_instruction_node.py | 85 +++++++++ 4 files changed, 260 insertions(+), 3 deletions(-) create mode 100644 nodes/fibo_edit_node.py create mode 100644 nodes/fibo_edit_structured_instruction_node.py diff --git a/__init__.py b/__init__.py index c401776..2f5aa16 100644 --- a/__init__.py +++ b/__init__.py @@ -39,7 +39,9 @@ from .nodes import ( VideoIncreaseResolutionNode, VideoEraseElementsNode, LoadVideoFramesNode, - PreviewVideoURLNode + PreviewVideoURLNode, + FIBOEditNode, + FIBOEditStructuredInstructionNode ) # Map the node class to a name used internally by ComfyUI @@ -84,7 +86,9 @@ NODE_CLASS_MAPPINGS = { "VideoIncreaseResolutionNode":VideoIncreaseResolutionNode, "VideoEraseElementsNode":VideoEraseElementsNode, "LoadVideoFramesNode":LoadVideoFramesNode, - "PreviewVideoURLNode":PreviewVideoURLNode + "PreviewVideoURLNode":PreviewVideoURLNode, + "FIBOEditNode": FIBOEditNode, + "FIBOEditStructuredInstructionNode": FIBOEditStructuredInstructionNode } # Map the node display name to the one shown in the ComfyUI node interface NODE_DISPLAY_NAME_MAPPINGS = { @@ -128,7 +132,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VideoIncreaseResolutionNode":"Bria Video Increase Resolution", "VideoEraseElementsNode":"Bria Video Erase Elements", "LoadVideoFramesNode":"Bria Load Video", - "PreviewVideoURLNode":"Bria Preview Video" + "PreviewVideoURLNode":"Bria Preview Video", + "FIBOEditNode": "FIBO - Edit", + "FIBOEditStructuredInstructionNode": "FIBO - Edit - Structured Instruction" } diff --git a/nodes/__init__.py b/nodes/__init__.py index 6b7cf8d..4cedb18 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -41,4 +41,7 @@ from .video_nodes.video_mask_by_prompt_node import VideoMaskByPromptNode from .video_nodes.video_mask_by_key_points_node import VideoMaskByKeyPointsNode from .video_nodes.load_video import LoadVideoFramesNode from .video_nodes.preview_video_node_from_url import PreviewVideoURLNode +from .fibo_edit_node import FIBOEditNode +from .fibo_edit_structured_instruction_node import FIBOEditStructuredInstructionNode + diff --git a/nodes/fibo_edit_node.py b/nodes/fibo_edit_node.py new file mode 100644 index 0000000..511d103 --- /dev/null +++ b/nodes/fibo_edit_node.py @@ -0,0 +1,163 @@ +import requests +import torch + +from .common import ( + deserialize_and_get_comfy_key, + postprocess_image, + preprocess_image, + image_to_base64, + poll_status_until_completed, + preprocess_mask, +) + + +class FIBOEditNode: + """FIBO Edit Node - Edit images with instructions""" + + api_url = "https://engine.prod.bria-api.com/v2/image/edit" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "api_token": ("STRING", {"default": "BRIA_API_TOKEN"}), + "instruction": ("STRING",), + "images": ("IMAGE",), + }, + "optional": { + "mask": ("MASK",), + "structured_instruction": ("STRING",), + "negative_prompt": ("STRING",), + "steps_num": ( + "INT", + { + "default": 50, + "min": 1, + "max": 100, + }, + ), + "guidance_scale": ( + "INT", + { + "default": 5, + "min": 1, + "max": 20, + }, + ), + "seed": ("INT", {"default": 123456}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "INT") + RETURN_NAMES = ("IMAGE", "structured_prompt", "seed") + CATEGORY = "API Nodes" + FUNCTION = "execute" + + def _validate_token(self, api_token: str): + if api_token.strip() == "" or api_token.strip() == "BRIA_API_TOKEN": + raise Exception("Please insert a valid API token.") + + def _build_payload( + self, + instruction, + images, + mask=None, + structured_instruction=None, + negative_prompt=None, + steps_num=50, + guidance_scale=5, + seed=123456, + ): + # Process images + if isinstance(images, torch.Tensor): + processed_images = preprocess_image(images) + else: + processed_images = images + + payload = { + "instruction": instruction, + "images": [image_to_base64(processed_images)], + "steps_num": steps_num, + "guidance_scale": guidance_scale, + "seed": seed, + } + + # Add optional mask + if mask is not None: + if isinstance(mask, torch.Tensor): + processed_mask = preprocess_mask(mask) + else: + processed_mask = mask + payload["mask"] = image_to_base64(processed_mask) + + # Add optional structured_instruction + if structured_instruction: + payload["structured_instruction"] = structured_instruction + + # Add optional negative_prompt + if negative_prompt: + payload["negative_prompt"] = negative_prompt + + return payload + + def execute( + self, + api_token, + instruction, + images, + mask=None, + structured_instruction=None, + negative_prompt=None, + steps_num=50, + guidance_scale=5, + seed=123456, + ): + self._validate_token(api_token) + payload = self._build_payload( + instruction, + images, + mask, + structured_instruction, + negative_prompt, + steps_num, + guidance_scale, + seed, + ) + api_token = deserialize_and_get_comfy_key(api_token) + + headers = {"Content-Type": "application/json", "api_token": api_token} + + try: + response = requests.post(self.api_url, json=payload, headers=headers) + + if response.status_code in (200, 202): + print( + f"Initial request successful to {self.api_url}, polling for completion..." + ) + response_dict = response.json() + 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_token) + + result = final_response.get("result", {}) + result_image_url = result.get("image_url") + structured_prompt = result.get("structured_prompt", "") + used_seed = result.get("seed") + + image_response = requests.get(result_image_url) + result_image = postprocess_image(image_response.content) + + return (result_image, structured_prompt, used_seed) + + raise Exception( + f"Error: API request failed with status code {response.status_code} {response.text}" + ) + + except Exception as e: + raise Exception(f"{e}") \ No newline at end of file diff --git a/nodes/fibo_edit_structured_instruction_node.py b/nodes/fibo_edit_structured_instruction_node.py new file mode 100644 index 0000000..fe65dc1 --- /dev/null +++ b/nodes/fibo_edit_structured_instruction_node.py @@ -0,0 +1,85 @@ +import requests +import torch + +from .common import ( + deserialize_and_get_comfy_key, + image_to_base64, + poll_status_until_completed, + preprocess_image, +) + + +class FIBOEditStructuredInstructionNode: + """FIBO Edit Structured Instruction Node - Generate structured instructions for image editing""" + + api_url = "https://engine.prod.bria-api.com/v2/structured_instruction/generate" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "api_token": ("STRING", {"default": "BRIA_API_TOKEN"}), + "image": ("IMAGE",), + "instruction": ("STRING",), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("structured_instruction",) + CATEGORY = "API Nodes" + FUNCTION = "execute" + + def _validate_token(self, api_token: str): + if api_token.strip() == "" or api_token.strip() == "BRIA_API_TOKEN": + raise Exception("Please insert a valid API token.") + + def _build_payload(self, image, instruction): + # Process image + if isinstance(image, torch.Tensor): + processed_image = preprocess_image(image) + else: + processed_image = image + + payload = { + "instruction": instruction, + "images": [image_to_base64(processed_image)], + } + + return payload + + def execute(self, api_token, image, instruction): + self._validate_token(api_token) + payload = self._build_payload(image, instruction) + api_token = deserialize_and_get_comfy_key(api_token) + + headers = {"Content-Type": "application/json", "api_token": api_token} + + try: + response = requests.post(self.api_url, json=payload, headers=headers) + + if response.status_code in (200, 202): + print( + f"Initial request successful to {self.api_url}, polling for completion..." + ) + response_dict = response.json() + 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_token) + + result = final_response.get("result", {}) + structured_instruction = result.get("structured_instruction", "") + + return (structured_instruction,) + + raise Exception( + f"Error: API request failed with status code {response.status_code} {response.text}" + ) + + except Exception as e: + raise Exception(f"{e}") \ No newline at end of file