Add a bunch of nodes, 3 ready to use, the rest waiting for endpoint support (#108)

This commit is contained in:
Jedrzej Kosinski 2025-05-02 05:49:05 -05:00 committed by Robin Huang
parent 491c020f45
commit 8daa075737
2 changed files with 88 additions and 0 deletions

View File

@ -127,6 +127,84 @@ def convert_image_to_base64(image: torch.Tensor):
return base64.b64encode(img_byte_arr.getvalue()).decode()
def convert_mask_to_image(mask: torch.Tensor):
"""
Make mask have the expected amount of dims (4) and channels (3) to be recognized as an image.
"""
mask = mask.unsqueeze(-1)
mask = torch.cat([mask]*3, dim=-1)
return mask
def handle_bfl_synchronous_operation(
operation: SynchronousOperation, timeout_bfl_calls=360
):
response_api: BFLFluxProGenerateResponse = operation.execute()
return _poll_until_generated(
response_api.polling_url, timeout=timeout_bfl_calls
)
def _poll_until_generated(polling_url: str, timeout=360):
# used bfl-comfy-nodes to verify code implementation:
# https://github.com/black-forest-labs/bfl-comfy-nodes/tree/main
start_time = time.time()
retries_404 = 0
max_retries_404 = 5
retry_404_seconds = 2
retry_202_seconds = 2
retry_pending_seconds = 1
request = requests.Request(method=HttpMethod.GET, url=polling_url)
# NOTE: should True loop be replaced with checking if workflow has been interrupted?
while True:
response = requests.Session().send(request.prepare())
if response.status_code == 200:
result = response.json()
if result["status"] == BFLStatus.ready:
img_url = result["result"]["sample"]
img_response = requests.get(img_url)
return process_image_response(img_response)
elif result["status"] in [
BFLStatus.request_moderated,
BFLStatus.content_moderated,
]:
status = result["status"]
raise Exception(
f"BFL API did not return an image due to: {status}."
)
elif result["status"] == BFLStatus.error:
raise Exception(f"BFL API encountered an error: {result}.")
elif result["status"] == BFLStatus.pending:
time.sleep(retry_pending_seconds)
continue
elif response.status_code == 404:
if retries_404 < max_retries_404:
retries_404 += 1
time.sleep(retry_404_seconds)
continue
raise Exception(
f"BFL API could not find task after {max_retries_404} tries."
)
elif response.status_code == 202:
time.sleep(retry_202_seconds)
elif time.time() - start_time > timeout:
raise Exception(
f"BFL API experienced a timeout; could not return request under {timeout} seconds."
)
else:
raise Exception(f"BFL API encountered an error: {response.json()}")
def convert_image_to_base64(image: torch.Tensor):
scaled_image = downscale_image_tensor(image, total_pixels=2048 * 2048)
# remove batch dimension if present
if len(scaled_image.shape) > 3:
scaled_image = scaled_image[0]
image_np = (scaled_image.numpy() * 255).astype(np.uint8)
img = Image.fromarray(image_np)
img_byte_arr = io.BytesIO()
img.save(img_byte_arr, format="PNG")
return base64.b64encode(img_byte_arr.getvalue()).decode()
class FluxProUltraImageNode(ComfyNodeABC):
"""
Generates images using Flux Pro 1.1 Ultra via api based on prompt and resolution.

View File

@ -116,6 +116,16 @@ def recraft_multipart_parser(data, parent_key=None, formatter: callable=None, co
if converted_to_check is None:
converted_to_check = []
def combine(self, other: SVG):
return SVG(self.data + other.data)
@staticmethod
def combine_all(svgs: list[SVG]):
all_svgs = []
for svg in svgs:
all_svgs.extend(svg.data)
return SVG(all_svgs)
if formatter is None:
formatter = lambda v: v # Multipart representation of value