From 00f9679bd1e2cefc2114e77208141d93babac487 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 3 May 2025 18:23:27 -0500 Subject: [PATCH] Add 8 nodes - 4 BFL, 4 Stability (#117) --- comfy_api_nodes/apinode_utils.py | 18 ++ comfy_api_nodes/apis/stability_api.py | 37 ++- comfy_api_nodes/nodes_bfl.py | 52 +++-- comfy_api_nodes/nodes_recraft.py | 10 +- comfy_api_nodes/nodes_stability.py | 319 +++++++++++++++++++++++++- 5 files changed, 399 insertions(+), 37 deletions(-) diff --git a/comfy_api_nodes/apinode_utils.py b/comfy_api_nodes/apinode_utils.py index baeaa0ceb..13f900a61 100644 --- a/comfy_api_nodes/apinode_utils.py +++ b/comfy_api_nodes/apinode_utils.py @@ -524,3 +524,21 @@ def upload_images_to_comfyapi( if idx_image >= batch_length: break return download_urls + + +def resize_mask_to_image(mask: torch.Tensor, image: torch.Tensor, + upscale_method="nearest-exact", crop="disabled", + allow_gradient=True, add_channel_dim=False): + """ + Resize mask to be the same dimensions as an image, while maintaining proper format for API calls. + """ + _, H, W, _ = image.shape + mask = mask.unsqueeze(-1) + mask = mask.movedim(-1,1) + mask = common_upscale(mask, width=W, height=H, upscale_method=upscale_method, crop=crop) + mask = mask.movedim(1,-1) + if not add_channel_dim: + mask = mask.squeeze(-1) + if not allow_gradient: + mask = (mask > 0.5).float() + return mask diff --git a/comfy_api_nodes/apis/stability_api.py b/comfy_api_nodes/apis/stability_api.py index d8d4f4e6a..47c87daec 100644 --- a/comfy_api_nodes/apis/stability_api.py +++ b/comfy_api_nodes/apis/stability_api.py @@ -53,8 +53,8 @@ class StabilityStylePreset(str, Enum): class Stability_SD3_5_Model(str, Enum): sd3_5_large = "sd3.5-large" - sd3_5_large_turbo = "sd3.5-large-turbo" - #sd3_5_medium = "sd3.5-medium" + # sd3_5_large_turbo = "sd3.5-large-turbo" + sd3_5_medium = "sd3.5-medium" class Stability_SD3_5_GenerationMode(str, Enum): @@ -76,6 +76,25 @@ class StabilityStable3_5Request(BaseModel): strength: Optional[confloat(ge=0.0, le=1.0)] = Field(None) +class StabilityUpscaleConservativeRequest(BaseModel): + prompt: str = Field(...) + negative_prompt: Optional[str] = Field(None) + seed: Optional[int] = Field(None) + output_format: Optional[str] = Field(StabilityFormat.png.value) + image: Optional[str] = Field(None) + creativity: Optional[confloat(ge=0.2, le=0.5)] = Field(None) + + +class StabilityUpscaleCreativeRequest(BaseModel): + prompt: str = Field(...) + negative_prompt: Optional[str] = Field(None) + seed: Optional[int] = Field(None) + output_format: Optional[str] = Field(StabilityFormat.png.value) + image: Optional[str] = Field(None) + creativity: Optional[confloat(ge=0.1, le=0.5)] = Field(None) + style_preset: Optional[str] = Field(None) + + class StabilityStableUltraRequest(BaseModel): prompt: str = Field(...) negative_prompt: Optional[str] = Field(None) @@ -92,3 +111,17 @@ class StabilityStableUltraResponse(BaseModel): finish_reason: Optional[str] = Field(None) seed: Optional[int] = Field(None) + +class StabilityResultsGetResponse(BaseModel): + image: Optional[str] = Field(None) + finish_reason: Optional[str] = Field(None) + seed: Optional[int] = Field(None) + id: Optional[str] = Field(None) + name: Optional[str] = Field(None) + errors: Optional[list[str]] = Field(None) + status: Optional[str] = Field(None) + result: Optional[str] = Field(None) + + +class StabilityAsyncResponse(BaseModel): + id: Optional[str] = Field(None) diff --git a/comfy_api_nodes/nodes_bfl.py b/comfy_api_nodes/nodes_bfl.py index fd58249c8..f1014459a 100644 --- a/comfy_api_nodes/nodes_bfl.py +++ b/comfy_api_nodes/nodes_bfl.py @@ -20,6 +20,7 @@ from comfy_api_nodes.apinode_utils import ( downscale_image_tensor, validate_aspect_ratio, process_image_response, + resize_mask_to_image, ) import numpy as np @@ -589,9 +590,11 @@ class FluxProFillNode(ComfyNodeABC): auth_token=None, **kwargs, ): + # prepare mask + mask = resize_mask_to_image(mask, image) + mask = convert_image_to_base64(convert_mask_to_image(mask)) # make sure image will have alpha channel removed image = convert_image_to_base64(image[:,:,:,:3]) - mask = convert_image_to_base64(convert_mask_to_image(mask)) operation = SynchronousOperation( endpoint=ApiEndpoint( @@ -641,20 +644,22 @@ class FluxProCannyNode(ComfyNodeABC): }, ), "canny_low_threshold": ( - IO.INT, + IO.FLOAT, { - "default": 0, - "min": 0, - "max": 500, + "default": 0.1, + "min": 0.01, + "max": 0.99, + "step": 0.01, "tooltip": "Low threshold for Canny edge detection; ignored if skip_processing is True" }, ), "canny_high_threshold": ( - IO.INT, + IO.FLOAT, { - "default": 0, - "min": 0, - "max": 500, + "default": 0.4, + "min": 0.01, + "max": 0.99, + "step": 0.01, "tooltip": "High threshold for Canny edge detection; ignored if skip_processing is True" }, ), @@ -712,8 +717,8 @@ class FluxProCannyNode(ComfyNodeABC): control_image: torch.Tensor, prompt: str, prompt_upsampling: bool, - canny_low_threshold: int, - canny_high_threshold: int, + canny_low_threshold: float, + canny_high_threshold: float, skip_preprocessing: bool, steps: int, guidance: float, @@ -724,6 +729,13 @@ class FluxProCannyNode(ComfyNodeABC): control_image = convert_image_to_base64(control_image[:,:,:,:3]) preprocessed_image = None + # scale canny threshold between 0-500, to match BFL's API + def scale_value(value: float, min_val=0, max_val=500): + return min_val + value * (max_val - min_val) + canny_low_threshold = int(round(scale_value(canny_low_threshold))) + canny_high_threshold = int(round(scale_value(canny_high_threshold))) + + if skip_preprocessing: preprocessed_image = control_image control_image = None @@ -849,7 +861,7 @@ class FluxProDepthNode(ComfyNodeABC): operation = SynchronousOperation( endpoint=ApiEndpoint( - path="/proxy/bfl/flux-pro-1.0-canny/generate", + path="/proxy/bfl/flux-pro-1.0-depth/generate", method=HttpMethod.POST, request_model=BFLFluxDepthImageRequest, response_model=BFLFluxProGenerateResponse, @@ -874,18 +886,18 @@ class FluxProDepthNode(ComfyNodeABC): NODE_CLASS_MAPPINGS = { "FluxProUltraImageNode": FluxProUltraImageNode, # "FluxProImageNode": FluxProImageNode, - # "FluxProExpandNode": FluxProExpandNode, - # "FluxProFillNode": FluxProFillNode, - # "FluxProCannyNode": FluxProCannyNode, - # "FluxProDepthNode": FluxProDepthNode, + "FluxProExpandNode": FluxProExpandNode, + "FluxProFillNode": FluxProFillNode, + "FluxProCannyNode": FluxProCannyNode, + "FluxProDepthNode": FluxProDepthNode, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "FluxProUltraImageNode": "Flux 1.1 [pro] Ultra Image", # "FluxProImageNode": "Flux 1.1 [pro] Image", - # "FluxProExpandNode": "Flux.1 Expand Image", - # "FluxProFillNode": "Flux.1 Fill Image", - # "FluxProCannyNode": "Flux.1 Canny Control Image", - # "FluxProDepthNode": "Flux.1 Depth Control Image", + "FluxProExpandNode": "Flux.1 Expand Image", + "FluxProFillNode": "Flux.1 Fill Image", + "FluxProCannyNode": "Flux.1 Canny Control Image", + "FluxProDepthNode": "Flux.1 Depth Control Image", } diff --git a/comfy_api_nodes/nodes_recraft.py b/comfy_api_nodes/nodes_recraft.py index 62cc503e2..eaa15f92b 100644 --- a/comfy_api_nodes/nodes_recraft.py +++ b/comfy_api_nodes/nodes_recraft.py @@ -1,6 +1,6 @@ from __future__ import annotations from inspect import cleandoc -from comfy.utils import ProgressBar, common_upscale +from comfy.utils import ProgressBar from comfy.comfy_types.node_typing import IO from comfy_api_nodes.apis.recraft_api import ( RecraftImageGenerationRequest, @@ -25,6 +25,7 @@ from comfy_api_nodes.apinode_utils import ( bytesio_to_image_tensor, download_url_to_bytesio, tensor_to_bytesio, + resize_mask_to_image, ) import folder_paths import json @@ -654,12 +655,7 @@ class RecraftImageInpaintingNode: ) # prepare mask tensor - _, H, W, _ = image.shape - mask = mask.unsqueeze(-1) - mask = mask.movedim(-1,1) - mask = common_upscale(mask, width=W, height=H, upscale_method="nearest-exact", crop="disabled") - mask = mask.movedim(1,-1) - mask = (mask > 0.5).float() + mask = resize_mask_to_image(mask, image, allow_gradient=False, add_channel_dim=True) images = [] total = image.shape[0] diff --git a/comfy_api_nodes/nodes_stability.py b/comfy_api_nodes/nodes_stability.py index 2ac6b30a2..50d264b5d 100644 --- a/comfy_api_nodes/nodes_stability.py +++ b/comfy_api_nodes/nodes_stability.py @@ -1,6 +1,10 @@ from inspect import cleandoc from comfy.comfy_types.node_typing import IO from comfy_api_nodes.apis.stability_api import ( + StabilityUpscaleConservativeRequest, + StabilityUpscaleCreativeRequest, + StabilityAsyncResponse, + StabilityResultsGetResponse, StabilityStable3_5Request, StabilityStableUltraRequest, StabilityStableUltraResponse, @@ -13,6 +17,8 @@ from comfy_api_nodes.apis.client import ( ApiEndpoint, HttpMethod, SynchronousOperation, + PollingOperation, + EmptyRequest, ) from comfy_api_nodes.apinode_utils import ( bytesio_to_image_tensor, @@ -22,8 +28,22 @@ from comfy_api_nodes.apinode_utils import ( import torch import base64 from io import BytesIO +from enum import Enum +class StabilityPollStatus(str, Enum): + finished = "finished" + in_progress = "in_progress" + failed = "failed" + + +def get_async_dummy_status(x: StabilityResultsGetResponse): + if x.name is not None or x.errors is not None: + return StabilityPollStatus.failed + elif x.finish_reason is not None: + return StabilityPollStatus.finished + return StabilityPollStatus.in_progress + class StabilityStableImageUltraNode: """ @@ -108,7 +128,7 @@ class StabilityStableImageUltraNode: # prepare image binary if image present image_binary = None if image is not None: - image_binary = tensor_to_bytesio(image, 1504 * 1504).read() + image_binary = tensor_to_bytesio(image, total_pixels=1504*1504).read() else: image_denoise = None @@ -174,6 +194,7 @@ class StabilityStableImageSD_3_5Node: "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results." }, ), + "model": ([x.value for x in Stability_SD3_5_Model],), "aspect_ratio": ([x.value for x in StabilityAspectRatio], { "default": StabilityAspectRatio.ratio_1_1, @@ -232,16 +253,16 @@ class StabilityStableImageSD_3_5Node: }, } - def api_call(self, prompt: str, aspect_ratio: str, style_preset: str, seed: int, cfg_scale: float, + def api_call(self, model: str, prompt: str, aspect_ratio: str, style_preset: str, seed: int, cfg_scale: float, negative_prompt: str=None, image: torch.Tensor = None, image_denoise: float=None, auth_token=None): - model = Stability_SD3_5_Model.sd3_5_large.value # prepare image binary if image present image_binary = None - mode = Stability_SD3_5_GenerationMode.text_to_image.value + mode = Stability_SD3_5_GenerationMode.text_to_image if image is not None: - image_binary = tensor_to_bytesio(image, 1504 * 1504).read() - mode = Stability_SD3_5_GenerationMode.image_to_image.value + image_binary = tensor_to_bytesio(image, total_pixels=1504*1504).read() + mode = Stability_SD3_5_GenerationMode.image_to_image + aspect_ratio = None else: image_denoise = None @@ -287,15 +308,297 @@ class StabilityStableImageSD_3_5Node: return (returned_image,) +class StabilityUpscaleConservativeNode: + """ + Upscale image with minimal alterations to 4K resolution. + """ + + RETURN_TYPES = (IO.IMAGE,) + DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value + FUNCTION = "api_call" + API_NODE = True + CATEGORY = "api node/image/stability" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": (IO.IMAGE,), + "prompt": ( + IO.STRING, + { + "multiline": True, + "default": "", + "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results." + }, + ), + "creativity": ( + IO.FLOAT, + { + "default": 0.35, + "min": 0.2, + "max": 0.5, + "step": 0.01, + "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init image.", + }, + ), + "seed": ( + IO.INT, + { + "default": 0, + "min": 0, + "max": 4294967294, + "control_after_generate": True, + "tooltip": "The random seed used for creating the noise.", + }, + ), + }, + "optional": { + "negative_prompt": ( + IO.STRING, + { + "default": "", + "forceInput": True, + "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature." + }, + ), + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + }, + } + + def api_call(self, image: torch.Tensor, prompt: str, creativity: float, seed: int, negative_prompt: str=None, + auth_token=None): + image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read() + + if not negative_prompt: + negative_prompt = None + + files = { + "image": image_binary + } + + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path="/proxy/stability/v2beta/stable-image/upscale/conservative", + method=HttpMethod.POST, + request_model=StabilityUpscaleConservativeRequest, + response_model=StabilityStableUltraResponse, + ), + request=StabilityUpscaleConservativeRequest( + prompt=prompt, + negative_prompt=negative_prompt, + creativity=round(creativity,2), + seed=seed, + ), + files=files, + content_type="multipart/form-data", + auth_token=auth_token, + ) + response_api = operation.execute() + + if response_api.finish_reason != "SUCCESS": + raise Exception(f"Stability Upscale Conservative generation failed: {response_api.finish_reason}.") + + image_data = base64.b64decode(response_api.image) + returned_image = bytesio_to_image_tensor(BytesIO(image_data)) + + return (returned_image,) + + +class StabilityUpscaleCreativeNode: + """ + Upscale image with minimal alterations to 4K resolution. + """ + + RETURN_TYPES = (IO.IMAGE,) + DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value + FUNCTION = "api_call" + API_NODE = True + CATEGORY = "api node/image/stability" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": (IO.IMAGE,), + "prompt": ( + IO.STRING, + { + "multiline": True, + "default": "", + "tooltip": "What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results." + }, + ), + "creativity": ( + IO.FLOAT, + { + "default": 0.3, + "min": 0.1, + "max": 0.5, + "step": 0.01, + "tooltip": "Controls the likelihood of creating additional details not heavily conditioned by the init image.", + }, + ), + "style_preset": (get_stability_style_presets(), + { + "tooltip": "Optional desired style of generated image.", + }, + ), + "seed": ( + IO.INT, + { + "default": 0, + "min": 0, + "max": 4294967294, + "control_after_generate": True, + "tooltip": "The random seed used for creating the noise.", + }, + ), + }, + "optional": { + "negative_prompt": ( + IO.STRING, + { + "default": "", + "forceInput": True, + "tooltip": "Keywords of what you do not wish to see in the output image. This is an advanced feature." + }, + ), + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + }, + } + + def api_call(self, image: torch.Tensor, prompt: str, creativity: float, style_preset: str, seed: int, negative_prompt: str=None, + auth_token=None): + image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read() + + if not negative_prompt: + negative_prompt = None + if style_preset == "None": + style_preset = None + + files = { + "image": image_binary + } + + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path="/proxy/stability/v2beta/stable-image/upscale/creative", + method=HttpMethod.POST, + request_model=StabilityUpscaleCreativeRequest, + response_model=StabilityAsyncResponse, + ), + request=StabilityUpscaleCreativeRequest( + prompt=prompt, + negative_prompt=negative_prompt, + creativity=round(creativity,2), + style_preset=style_preset, + seed=seed, + ), + files=files, + content_type="multipart/form-data", + auth_token=auth_token, + ) + response_api = operation.execute() + + operation = PollingOperation( + poll_endpoint=ApiEndpoint( + path=f"/proxy/stability/v2beta/results/{response_api.id}", + method=HttpMethod.GET, + request_model=EmptyRequest, + response_model=StabilityResultsGetResponse, + ), + poll_interval=3, + completed_statuses=[StabilityPollStatus.finished], + failed_statuses=[StabilityPollStatus.failed], + status_extractor=lambda x: get_async_dummy_status(x), + auth_token=auth_token, + ) + response_poll: StabilityResultsGetResponse = operation.execute() + + if response_poll.finish_reason != "SUCCESS": + raise Exception(f"Stability Upscale Creative generation failed: {response_poll.finish_reason}.") + + image_data = base64.b64decode(response_poll.result) + returned_image = bytesio_to_image_tensor(BytesIO(image_data)) + + return (returned_image,) + + +class StabilityUpscaleFastNode: + """ + Quickly upscales an image via Stability API call to 4x its original size; intended for upscaling low-quality/compressed images. + """ + + RETURN_TYPES = (IO.IMAGE,) + DESCRIPTION = cleandoc(__doc__ or "") # Handle potential None value + FUNCTION = "api_call" + API_NODE = True + CATEGORY = "api node/image/stability" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": (IO.IMAGE,), + }, + "optional": { + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + }, + } + + def api_call(self, image: torch.Tensor, + auth_token=None): + image_binary = tensor_to_bytesio(image, total_pixels=4096*4096).read() + + files = { + "image": image_binary + } + + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path="/proxy/stability/v2beta/stable-image/upscale/fast", + method=HttpMethod.POST, + request_model=EmptyRequest, + response_model=StabilityStableUltraResponse, + ), + request=EmptyRequest(), + files=files, + content_type="multipart/form-data", + auth_token=auth_token, + ) + response_api = operation.execute() + + if response_api.finish_reason != "SUCCESS": + raise Exception(f"Stability Upscale Fast failed: {response_api.finish_reason}.") + + image_data = base64.b64decode(response_api.image) + returned_image = bytesio_to_image_tensor(BytesIO(image_data)) + + return (returned_image,) + + # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "StabilityStableImageUltraNode": StabilityStableImageUltraNode, - # "StabilityStableImageSD_3_5Node": StabilityStableImageSD_3_5Node, + "StabilityStableImageSD_3_5Node": StabilityStableImageSD_3_5Node, + "StabilityUpscaleConservativeNode": StabilityUpscaleConservativeNode, + "StabilityUpscaleCreativeNode": StabilityUpscaleCreativeNode, + "StabilityUpscaleFastNode": StabilityUpscaleFastNode, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "StabilityStableImageUltraNode": "Stability Stable Image Ultra", - # "StabilityStableImageSD_3_5Node": "Stability Stable Diffusion 3.5 Image", + "StabilityStableImageSD_3_5Node": "Stability Stable Diffusion 3.5 Image", + "StabilityUpscaleConservativeNode": "Stability Upscale Conservative", + "StabilityUpscaleCreativeNode": "Stability Upscale Creative", + "StabilityUpscaleFastNode": "Stability Upscale Fast", }