Merge pull request #33 from Bria-AI/WAI-4234

WAI-4234
This commit is contained in:
Yazan Numoor
2025-12-16 08:27:11 +02:00
committed by GitHub
8 changed files with 1424 additions and 404 deletions
+14
View File
@@ -15,7 +15,11 @@ from .nodes import (
TailoredPortraitNode,
ReimagineNode,
GenerateImageNodeV2,
GenerateImageLiteNodeV2,
RefineImageNodeV2,
RefineImageLiteNodeV2,
GenerateStructuredPromptNodeV2,
GenerateStructuredPromptLiteNodeV2,
ShotByTextAutomaticNode,
ShotByImageManualPaddingNode,
ShotByImageAutomaticAspectRatioNode,
@@ -58,7 +62,11 @@ NODE_CLASS_MAPPINGS = {
"ReimagineNode": ReimagineNode,
"AttributionByImageNode": AttributionByImageNode,
"GenerateImageNodeV2": GenerateImageNodeV2,
"GenerateImageLiteNodeV2": GenerateImageLiteNodeV2,
"RefineImageNodeV2": RefineImageNodeV2,
"RefineImageLiteNodeV2": RefineImageLiteNodeV2,
"GenerateStructuredPromptNodeV2": GenerateStructuredPromptNodeV2,
"GenerateStructuredPromptLiteNodeV2": GenerateStructuredPromptLiteNodeV2,
}
# Map the node display name to the one shown in the ComfyUI node interface
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -89,5 +97,11 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"ReimagineNode": "Bria Reimagine",
"AttributionByImageNode": "Attribution By Image Node",
"GenerateImageNodeV2": "Generate Image",
"GenerateImageLiteNodeV2": "Generate Image - Lite",
"RefineImageNodeV2": "Refine and Regenerate Image",
"RefineImageLiteNodeV2": "Refine Image - Lite",
"GenerateStructuredPromptNodeV2": "Generate Structured Prompt",
"GenerateStructuredPromptLiteNodeV2": "Generate Structured Prompt - Lite",
}
+6 -2
View File
@@ -11,8 +11,12 @@ 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
from .generate_image_node_v2 import GenerateImageNodeV2
from .refine_image_node_v2 import RefineImageNodeV2
from .generate_image_node_v2 import GenerateImageNodeV2, GenerateImageLiteNodeV2
from .refine_image_node_v2 import RefineImageNodeV2, RefineImageLiteNodeV2
from .generate_structured_prompt_node_v2 import (
GenerateStructuredPromptNodeV2,
GenerateStructuredPromptLiteNodeV2,
)
from .shot_by_text_node import ShotByTextOriginalNode
from .shot_by_text_automatic_aspect_ratio_node import ShotByTextAutomaticAspectRatioNode
from .shot_by_text_automatic_node import ShotByTextAutomaticNode
+50 -17
View File
@@ -14,26 +14,44 @@ class _BaseGenerateImageNodeV2:
"""Base class for image generation nodes (standard & pro)."""
api_url = None # Each subclass must define its API endpoint
supports_negative_prompt = True # Can be overridden by subclasses
# STEP SETTINGS (defaults for standard API)
default_steps = 50
min_steps = 35
max_steps = 50
@classmethod
def INPUT_TYPES(cls):
optional_inputs = {
"model_version": (["FIBO"], {"default": "FIBO"}),
"structured_prompt": ("STRING",),
"images": ("IMAGE",),
"aspect_ratio": (
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"],
{"default": "1:1"},
),
"steps_num": (
"INT",
{
"default": cls.default_steps,
"min": cls.min_steps,
"max": cls.max_steps,
},
),
"guidance_scale": ("INT", {"default": 5, "min": 3, "max": 5}),
"seed": ("INT", {"default": 123456}),
}
# Add negative_prompt only if supported
if cls.supports_negative_prompt:
optional_inputs["negative_prompt"] = ("STRING", {"default": ""})
return {
"required": {
"api_token": ("STRING", {"default": "BRIA_API_TOKEN"}),
"prompt": ("STRING",),
},
"optional": {
"model_version": (["FIBO"], {"default": "FIBO"}),
"negative_prompt": ("STRING", {"default": ""}),
"images": ("IMAGE",),
"aspect_ratio": (
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"],
{"default": "1:1"},
),
"steps_num": ("INT", {"default": 50, "min": 20, "max": 50}),
"guidance_scale": ("INT", {"default": 5, "min": 3, "max": 5}),
"seed": ("INT", {"default": 123456}),
},
"optional": optional_inputs,
}
RETURN_TYPES = ("IMAGE", "STRING", "INT")
@@ -49,29 +67,33 @@ class _BaseGenerateImageNodeV2:
self,
prompt,
model_version,
negative_prompt,
structured_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt=None,
images=None,
):
payload = {
"prompt": prompt,
"model_version": model_version,
"negative_prompt": negative_prompt,
"structured_prompt":structured_prompt,
"aspect_ratio": aspect_ratio,
"steps_num": steps_num,
"guidance_scale": guidance_scale,
"seed": seed,
}
# Add negative_prompt only if supported and provided
if self.supports_negative_prompt and negative_prompt is not None:
payload["negative_prompt"] = negative_prompt
if images is not None:
if isinstance(images, torch.Tensor):
preprocess_images = preprocess_image(images)
payload["images"] = [image_to_base64(preprocess_images)]
return payload
def execute(
@@ -79,22 +101,24 @@ class _BaseGenerateImageNodeV2:
api_token,
prompt,
model_version,
negative_prompt,
structured_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt=None,
images=None,
):
self._validate_token(api_token)
payload = self._build_payload(
prompt,
model_version,
negative_prompt,
structured_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt,
images,
)
api_token = deserialize_and_get_comfy_key(api_token)
@@ -141,3 +165,12 @@ class GenerateImageNodeV2(_BaseGenerateImageNodeV2):
"""Standard Image Generation Node"""
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/image/generate"
class GenerateImageLiteNodeV2(_BaseGenerateImageNodeV2):
"""Lite Image Generation Node"""
supports_negative_prompt = False
default_steps = 8
min_steps = 8
max_steps = 30
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/image/generate/lite"
+118
View File
@@ -0,0 +1,118 @@
import requests
from .common import (
deserialize_and_get_comfy_key,
image_to_base64,
poll_status_until_completed,
preprocess_image,
)
import torch
class _BaseGenerateStructuredPromptNodeV2:
"""Base class for structured prompt generation nodes (standard & lite)."""
api_url = None # Each subclass must define its API endpoint
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"api_token": ("STRING", {"default": "BRIA_API_TOKEN"}),
"prompt": ("STRING",),
},
"optional": {
"structured_prompt": ("STRING",),
"images": ("IMAGE",),
"seed": ("INT", {"default": 123456}),
},
}
RETURN_TYPES = ("STRING", "INT")
RETURN_NAMES = ("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,
prompt,
seed,
structured_prompt,
images=None
):
payload = {
"prompt": prompt,
"seed": seed,
"structured_prompt":structured_prompt
}
if images is not None:
if isinstance(images, torch.Tensor):
preprocess_images = preprocess_image(images)
payload["images"] = [image_to_base64(preprocess_images)]
return payload
def execute(
self,
api_token,
prompt,
seed,
structured_prompt,
images=None,
):
self._validate_token(api_token)
payload = self._build_payload(
prompt,
seed,
structured_prompt,
images
)
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_prompt = result.get("structured_prompt", "")
used_seed = result.get("seed", seed)
return (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}")
class GenerateStructuredPromptNodeV2(_BaseGenerateStructuredPromptNodeV2):
"""Standard Structured Prompt Generation Node"""
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate"
class GenerateStructuredPromptLiteNodeV2(_BaseGenerateStructuredPromptNodeV2):
"""Lite Structured Prompt Generation Node"""
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate/lite"
+56 -19
View File
@@ -7,26 +7,42 @@ class _BaseRefineImageNodeV2:
api_url = None # Must be overridden by subclasses
generate_api_url = None
supports_negative_prompt = True # Can be overridden by subclasses
# STEP SETTINGS (defaults for standard API)
default_steps = 50
min_steps = 35
max_steps = 50
@classmethod
def INPUT_TYPES(cls):
optional_inputs = {
"model_version": (["FIBO"], {"default": "FIBO"}),
"aspect_ratio": (
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"],
{"default": "1:1"},
),
"steps_num": (
"INT",
{
"default": cls.default_steps,
"min": cls.min_steps,
"max": cls.max_steps,
},
),
"guidance_scale": ("INT", {"default": 5, "min": 3, "max": 5}),
"seed": ("INT", {"default": 123456}),
}
# Add negative_prompt only if supported
if cls.supports_negative_prompt:
optional_inputs["negative_prompt"] = ("STRING", {"default": ""})
return {
"required": {
"api_token": ("STRING", {"default": "BRIA_API_TOKEN"}),
"prompt": ("STRING",),
"structured_prompt": ("STRING",),
},
"optional": {
"model_version": (["FIBO"], {"default": "FIBO"}),
"negative_prompt": ("STRING", {"default": ""}),
"aspect_ratio": (
["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"],
{"default": "1:1"},
),
"steps_num": ("INT", {"default": 50, "min": 20, "max": 50}),
"guidance_scale": ("INT", {"default": 5, "min": 3, "max": 5}),
"seed": ("INT", {"default": 123456}),
},
"optional": optional_inputs,
}
RETURN_TYPES = ("IMAGE", "STRING", "INT")
@@ -43,22 +59,26 @@ class _BaseRefineImageNodeV2:
prompt,
structured_prompt,
model_version,
negative_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt=None,
):
return {
payload = {
"prompt": prompt,
"model_version": model_version,
"negative_prompt": negative_prompt,
"aspect_ratio": aspect_ratio,
"steps_num": steps_num,
"guidance_scale": guidance_scale,
"seed": seed,
"structured_prompt": structured_prompt,
}
if self.supports_negative_prompt and negative_prompt is not None:
payload["negative_prompt"] = negative_prompt
return payload
def execute(
self,
@@ -66,22 +86,22 @@ class _BaseRefineImageNodeV2:
prompt,
structured_prompt,
model_version,
negative_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt=None,
):
self._validate_token(api_token)
payload = self._build_payload(
prompt,
structured_prompt,
model_version,
negative_prompt,
aspect_ratio,
steps_num,
guidance_scale,
seed,
negative_prompt,
)
api_token = deserialize_and_get_comfy_key(api_token)
headers = {"Content-Type": "application/json", "api_token": api_token}
@@ -111,12 +131,16 @@ class _BaseRefineImageNodeV2:
"prompt": prompt,
"structured_prompt":structured_prompt,
"model_version": model_version,
"negative_prompt": negative_prompt,
"aspect_ratio": aspect_ratio,
"steps_num": steps_num,
"guidance_scale": guidance_scale,
"seed": used_seed,
}
# Add negative_prompt only if supported and provided
if self.supports_negative_prompt and negative_prompt is not None:
payloadForImageGenetrate["negative_prompt"] = negative_prompt
headers = {"Content-Type": "application/json", "api_token": api_token}
response = requests.post(self.generate_api_url, json=payloadForImageGenetrate, headers=headers)
@@ -158,4 +182,17 @@ class RefineImageNodeV2(_BaseRefineImageNodeV2):
"""Standard Refine Image Node"""
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate"
self.generate_api_url = "https://engine.prod.bria-api.com/v2/image/generate"
self.generate_api_url = "https://engine.prod.bria-api.com/v2/image/generate"
class RefineImageLiteNodeV2(_BaseRefineImageNodeV2):
"""Lite Refine Image Node"""
supports_negative_prompt = False
supports_negative_prompt = False
default_steps = 8
min_steps = 8
max_steps = 30
def __init__(self):
self.api_url = "https://engine.prod.bria-api.com/v2/structured_prompt/generate/lite"
self.generate_api_url = "https://engine.prod.bria-api.com/v2/image/generate/lite"
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui-bria-api"
description = "Custom nodes for ComfyUI using BRIA's API."
version = "2.1.8"
version = "2.1.9"
license = {file = "LICENSE"}
[project.urls]
@@ -0,0 +1,661 @@
{
"id": "1c31a92d-1e48-46b0-b4ad-60d537260922",
"revision": 0,
"last_node_id": 27,
"last_link_id": 46,
"nodes": [
{
"id": 5,
"type": "PreviewImage",
"pos": [
1235.1895751953125,
10.958272933959961
],
"size": [
140,
26
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 43
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 6,
"type": "PreviewImage",
"pos": [
774.4347534179688,
-58.6872444152832
],
"size": [
140,
26
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 29
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 1,
"type": "Note",
"pos": [
-175.84976196289062,
102.76270294189453
],
"size": [
210,
88
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {},
"widgets_values": [
"If you would like to start with prompt"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 2,
"type": "Note",
"pos": [
-225.1866455078125,
624.7235107421875
],
"size": [
210,
88
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {},
"widgets_values": [
"If you would like to start with reference image + prompt"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 15,
"type": "PreviewImage",
"pos": [
719.5359497070312,
443.11285400390625
],
"size": [
140,
26
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 36
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 17,
"type": "PreviewImage",
"pos": [
1132.783935546875,
517.3931884765625
],
"size": [
140,
26
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 44
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 20,
"type": "GenerateStructuredPromptLiteNodeV2",
"pos": [
68.9760971069336,
82.11421966552734
],
"size": [
310.7749938964844,
174
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "images",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "structured_prompt",
"type": "STRING",
"links": [
27
]
},
{
"name": "seed",
"type": "INT",
"links": [
28
]
}
],
"properties": {
"Node name for S&R": "GenerateStructuredPromptLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"",
123456,
"randomize"
]
},
{
"id": 24,
"type": "GenerateImageLiteNodeV2",
"pos": [
460.6510314941406,
44.938873291015625
],
"size": [
270,
290
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "images",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "structured_prompt",
"shape": 7,
"type": "STRING",
"widget": {
"name": "structured_prompt"
},
"link": 27
},
{
"name": "seed",
"shape": 7,
"type": "INT",
"widget": {
"name": "seed"
},
"link": 28
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
29
]
},
{
"name": "structured_prompt",
"type": "STRING",
"links": [
41
]
},
{
"name": "seed",
"type": "INT",
"links": [
42
]
}
],
"properties": {
"Node name for S&R": "GenerateImageLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"FIBO",
"",
"1:1",
8,
5,
123456,
"randomize"
]
},
{
"id": 26,
"type": "RefineImageLiteNodeV2",
"pos": [
834.9641723632812,
46.47432327270508
],
"size": [
270,
290
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "structured_prompt",
"type": "STRING",
"widget": {
"name": "structured_prompt"
},
"link": 41
},
{
"name": "seed",
"shape": 7,
"type": "INT",
"widget": {
"name": "seed"
},
"link": 42
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
43
]
},
{
"name": "structured_prompt",
"type": "STRING",
"links": null
},
{
"name": "seed",
"type": "INT",
"links": null
}
],
"properties": {
"Node name for S&R": "RefineImageLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"",
"FIBO",
"1:1",
8,
5,
123456,
"randomize"
]
},
{
"id": 25,
"type": "GenerateImageLiteNodeV2",
"pos": [
404.58843994140625,
532.5309448242188
],
"size": [
270,
290
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "images",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "structured_prompt",
"shape": 7,
"type": "STRING",
"widget": {
"name": "structured_prompt"
},
"link": 34
},
{
"name": "seed",
"shape": 7,
"type": "INT",
"widget": {
"name": "seed"
},
"link": 35
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
36
]
},
{
"name": "structured_prompt",
"type": "STRING",
"links": [
45
]
},
{
"name": "seed",
"type": "INT",
"links": [
46
]
}
],
"properties": {
"Node name for S&R": "GenerateImageLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"FIBO",
"",
"1:1",
8,
5,
123456,
"randomize"
]
},
{
"id": 27,
"type": "RefineImageLiteNodeV2",
"pos": [
820.26611328125,
558.9785766601562
],
"size": [
270,
290
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "structured_prompt",
"type": "STRING",
"widget": {
"name": "structured_prompt"
},
"link": 45
},
{
"name": "seed",
"shape": 7,
"type": "INT",
"widget": {
"name": "seed"
},
"link": 46
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
44
]
},
{
"name": "structured_prompt",
"type": "STRING",
"links": null
},
{
"name": "seed",
"type": "INT",
"links": null
}
],
"properties": {
"Node name for S&R": "RefineImageLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"",
"FIBO",
"1:1",
8,
5,
123456,
"randomize"
]
},
{
"id": 21,
"type": "GenerateStructuredPromptLiteNodeV2",
"pos": [
3.162916421890259,
587.301025390625
],
"size": [
310.7749938964844,
174
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "structured_prompt",
"type": "STRING",
"links": [
34
]
},
{
"name": "seed",
"type": "INT",
"links": [
35
]
}
],
"properties": {
"Node name for S&R": "GenerateStructuredPromptLiteNodeV2"
},
"widgets_values": [
"BRIA_API_TOKEN",
"",
"",
123456,
"randomize"
]
}
],
"links": [
[
27,
20,
0,
24,
1,
"STRING"
],
[
28,
20,
1,
24,
2,
"INT"
],
[
29,
24,
0,
6,
0,
"IMAGE"
],
[
34,
21,
0,
25,
1,
"STRING"
],
[
35,
21,
1,
25,
2,
"INT"
],
[
36,
25,
0,
15,
0,
"IMAGE"
],
[
41,
24,
1,
26,
0,
"STRING"
],
[
42,
24,
2,
26,
1,
"INT"
],
[
43,
26,
0,
5,
0,
"IMAGE"
],
[
44,
27,
0,
17,
0,
"IMAGE"
],
[
45,
25,
1,
27,
0,
"STRING"
],
[
46,
25,
2,
27,
1,
"INT"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.6276708501927047,
"offset": [
650.5506889681789,
132.55696596915172
]
},
"frontendVersion": "1.25.11"
},
"version": 0.4
}
File diff suppressed because it is too large Load Diff