mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2025-12-08 21:44:33 +08:00
658 lines
22 KiB
Python
658 lines
22 KiB
Python
from inspect import cleandoc
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from typing_extensions import override
|
|
|
|
from comfy_api.latest import IO, ComfyExtension
|
|
from comfy_api_nodes.apis.bfl_api import (
|
|
BFLFluxExpandImageRequest,
|
|
BFLFluxFillImageRequest,
|
|
BFLFluxKontextProGenerateRequest,
|
|
BFLFluxProGenerateRequest,
|
|
BFLFluxProGenerateResponse,
|
|
BFLFluxProUltraGenerateRequest,
|
|
BFLFluxStatusResponse,
|
|
BFLStatus,
|
|
)
|
|
from comfy_api_nodes.util import (
|
|
ApiEndpoint,
|
|
download_url_to_image_tensor,
|
|
poll_op,
|
|
resize_mask_to_image,
|
|
sync_op,
|
|
tensor_to_base64_string,
|
|
validate_aspect_ratio_string,
|
|
validate_string,
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
class FluxProUltraImageNode(IO.ComfyNode):
|
|
"""
|
|
Generates images using Flux Pro 1.1 Ultra via api based on prompt and resolution.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="FluxProUltraImageNode",
|
|
display_name="Flux 1.1 [pro] Ultra Image",
|
|
category="api node/image/BFL",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Prompt for the image generation",
|
|
),
|
|
IO.Boolean.Input(
|
|
"prompt_upsampling",
|
|
default=False,
|
|
tooltip="Whether to perform upsampling on the prompt. "
|
|
"If active, automatically modifies the prompt for more creative generation, "
|
|
"but results are nondeterministic (same seed will not produce exactly the same result).",
|
|
),
|
|
IO.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip="The random seed used for creating the noise.",
|
|
),
|
|
IO.String.Input(
|
|
"aspect_ratio",
|
|
default="16:9",
|
|
tooltip="Aspect ratio of image; must be between 1:4 and 4:1.",
|
|
),
|
|
IO.Boolean.Input(
|
|
"raw",
|
|
default=False,
|
|
tooltip="When True, generate less processed, more natural-looking images.",
|
|
),
|
|
IO.Image.Input(
|
|
"image_prompt",
|
|
optional=True,
|
|
),
|
|
IO.Float.Input(
|
|
"image_prompt_strength",
|
|
default=0.1,
|
|
min=0.0,
|
|
max=1.0,
|
|
step=0.01,
|
|
tooltip="Blend between the prompt and the image prompt.",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[IO.Image.Output()],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, aspect_ratio: str):
|
|
validate_aspect_ratio_string(aspect_ratio, (1, 4), (4, 1))
|
|
return True
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
prompt: str,
|
|
aspect_ratio: str,
|
|
prompt_upsampling: bool = False,
|
|
raw: bool = False,
|
|
seed: int = 0,
|
|
image_prompt: Optional[torch.Tensor] = None,
|
|
image_prompt_strength: float = 0.1,
|
|
) -> IO.NodeOutput:
|
|
if image_prompt is None:
|
|
validate_string(prompt, strip_whitespace=False)
|
|
initial_response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/bfl/flux-pro-1.1-ultra/generate", method="POST"),
|
|
response_model=BFLFluxProGenerateResponse,
|
|
data=BFLFluxProUltraGenerateRequest(
|
|
prompt=prompt,
|
|
prompt_upsampling=prompt_upsampling,
|
|
seed=seed,
|
|
aspect_ratio=aspect_ratio,
|
|
raw=raw,
|
|
image_prompt=(image_prompt if image_prompt is None else tensor_to_base64_string(image_prompt)),
|
|
image_prompt_strength=(None if image_prompt is None else round(image_prompt_strength, 2)),
|
|
),
|
|
)
|
|
response = await poll_op(
|
|
cls,
|
|
ApiEndpoint(initial_response.polling_url),
|
|
response_model=BFLFluxStatusResponse,
|
|
status_extractor=lambda r: r.status,
|
|
progress_extractor=lambda r: r.progress,
|
|
completed_statuses=[BFLStatus.ready],
|
|
failed_statuses=[
|
|
BFLStatus.request_moderated,
|
|
BFLStatus.content_moderated,
|
|
BFLStatus.error,
|
|
BFLStatus.task_not_found,
|
|
],
|
|
queued_statuses=[],
|
|
)
|
|
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
|
|
|
|
|
class FluxKontextProImageNode(IO.ComfyNode):
|
|
"""
|
|
Edits images using Flux.1 Kontext [pro] via api based on prompt and aspect ratio.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id=cls.NODE_ID,
|
|
display_name=cls.DISPLAY_NAME,
|
|
category="api node/image/BFL",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Prompt for the image generation - specify what and how to edit.",
|
|
),
|
|
IO.String.Input(
|
|
"aspect_ratio",
|
|
default="16:9",
|
|
tooltip="Aspect ratio of image; must be between 1:4 and 4:1.",
|
|
),
|
|
IO.Float.Input(
|
|
"guidance",
|
|
default=3.0,
|
|
min=0.1,
|
|
max=99.0,
|
|
step=0.1,
|
|
tooltip="Guidance strength for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"steps",
|
|
default=50,
|
|
min=1,
|
|
max=150,
|
|
tooltip="Number of steps for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"seed",
|
|
default=1234,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip="The random seed used for creating the noise.",
|
|
),
|
|
IO.Boolean.Input(
|
|
"prompt_upsampling",
|
|
default=False,
|
|
tooltip="Whether to perform upsampling on the prompt. If active, automatically modifies the prompt for more creative generation, but results are nondeterministic (same seed will not produce exactly the same result).",
|
|
),
|
|
IO.Image.Input(
|
|
"input_image",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[IO.Image.Output()],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
BFL_PATH = "/proxy/bfl/flux-kontext-pro/generate"
|
|
NODE_ID = "FluxKontextProImageNode"
|
|
DISPLAY_NAME = "Flux.1 Kontext [pro] Image"
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
prompt: str,
|
|
aspect_ratio: str,
|
|
guidance: float,
|
|
steps: int,
|
|
input_image: Optional[torch.Tensor] = None,
|
|
seed=0,
|
|
prompt_upsampling=False,
|
|
) -> IO.NodeOutput:
|
|
validate_aspect_ratio_string(aspect_ratio, (1, 4), (4, 1))
|
|
if input_image is None:
|
|
validate_string(prompt, strip_whitespace=False)
|
|
initial_response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(path=cls.BFL_PATH, method="POST"),
|
|
response_model=BFLFluxProGenerateResponse,
|
|
data=BFLFluxKontextProGenerateRequest(
|
|
prompt=prompt,
|
|
prompt_upsampling=prompt_upsampling,
|
|
guidance=round(guidance, 1),
|
|
steps=steps,
|
|
seed=seed,
|
|
aspect_ratio=aspect_ratio,
|
|
input_image=(input_image if input_image is None else tensor_to_base64_string(input_image)),
|
|
),
|
|
)
|
|
response = await poll_op(
|
|
cls,
|
|
ApiEndpoint(initial_response.polling_url),
|
|
response_model=BFLFluxStatusResponse,
|
|
status_extractor=lambda r: r.status,
|
|
progress_extractor=lambda r: r.progress,
|
|
completed_statuses=[BFLStatus.ready],
|
|
failed_statuses=[
|
|
BFLStatus.request_moderated,
|
|
BFLStatus.content_moderated,
|
|
BFLStatus.error,
|
|
BFLStatus.task_not_found,
|
|
],
|
|
queued_statuses=[],
|
|
)
|
|
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
|
|
|
|
|
class FluxKontextMaxImageNode(FluxKontextProImageNode):
|
|
"""
|
|
Edits images using Flux.1 Kontext [max] via api based on prompt and aspect ratio.
|
|
"""
|
|
|
|
DESCRIPTION = cleandoc(__doc__ or "")
|
|
BFL_PATH = "/proxy/bfl/flux-kontext-max/generate"
|
|
NODE_ID = "FluxKontextMaxImageNode"
|
|
DISPLAY_NAME = "Flux.1 Kontext [max] Image"
|
|
|
|
|
|
class FluxProImageNode(IO.ComfyNode):
|
|
"""
|
|
Generates images synchronously based on prompt and resolution.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="FluxProImageNode",
|
|
display_name="Flux 1.1 [pro] Image",
|
|
category="api node/image/BFL",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Prompt for the image generation",
|
|
),
|
|
IO.Boolean.Input(
|
|
"prompt_upsampling",
|
|
default=False,
|
|
tooltip="Whether to perform upsampling on the prompt. "
|
|
"If active, automatically modifies the prompt for more creative generation, "
|
|
"but results are nondeterministic (same seed will not produce exactly the same result).",
|
|
),
|
|
IO.Int.Input(
|
|
"width",
|
|
default=1024,
|
|
min=256,
|
|
max=1440,
|
|
step=32,
|
|
),
|
|
IO.Int.Input(
|
|
"height",
|
|
default=768,
|
|
min=256,
|
|
max=1440,
|
|
step=32,
|
|
),
|
|
IO.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip="The random seed used for creating the noise.",
|
|
),
|
|
IO.Image.Input(
|
|
"image_prompt",
|
|
optional=True,
|
|
),
|
|
# "image_prompt_strength": (
|
|
# IO.FLOAT,
|
|
# {
|
|
# "default": 0.1,
|
|
# "min": 0.0,
|
|
# "max": 1.0,
|
|
# "step": 0.01,
|
|
# "tooltip": "Blend between the prompt and the image prompt.",
|
|
# },
|
|
# ),
|
|
],
|
|
outputs=[IO.Image.Output()],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
prompt: str,
|
|
prompt_upsampling,
|
|
width: int,
|
|
height: int,
|
|
seed=0,
|
|
image_prompt=None,
|
|
# image_prompt_strength=0.1,
|
|
) -> IO.NodeOutput:
|
|
image_prompt = image_prompt if image_prompt is None else tensor_to_base64_string(image_prompt)
|
|
initial_response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(
|
|
path="/proxy/bfl/flux-pro-1.1/generate",
|
|
method="POST",
|
|
),
|
|
response_model=BFLFluxProGenerateResponse,
|
|
data=BFLFluxProGenerateRequest(
|
|
prompt=prompt,
|
|
prompt_upsampling=prompt_upsampling,
|
|
width=width,
|
|
height=height,
|
|
seed=seed,
|
|
image_prompt=image_prompt,
|
|
),
|
|
)
|
|
response = await poll_op(
|
|
cls,
|
|
ApiEndpoint(initial_response.polling_url),
|
|
response_model=BFLFluxStatusResponse,
|
|
status_extractor=lambda r: r.status,
|
|
progress_extractor=lambda r: r.progress,
|
|
completed_statuses=[BFLStatus.ready],
|
|
failed_statuses=[
|
|
BFLStatus.request_moderated,
|
|
BFLStatus.content_moderated,
|
|
BFLStatus.error,
|
|
BFLStatus.task_not_found,
|
|
],
|
|
queued_statuses=[],
|
|
)
|
|
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
|
|
|
|
|
class FluxProExpandNode(IO.ComfyNode):
|
|
"""
|
|
Outpaints image based on prompt.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="FluxProExpandNode",
|
|
display_name="Flux.1 Expand Image",
|
|
category="api node/image/BFL",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("image"),
|
|
IO.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Prompt for the image generation",
|
|
),
|
|
IO.Boolean.Input(
|
|
"prompt_upsampling",
|
|
default=False,
|
|
tooltip="Whether to perform upsampling on the prompt. "
|
|
"If active, automatically modifies the prompt for more creative generation, "
|
|
"but results are nondeterministic (same seed will not produce exactly the same result).",
|
|
),
|
|
IO.Int.Input(
|
|
"top",
|
|
default=0,
|
|
min=0,
|
|
max=2048,
|
|
tooltip="Number of pixels to expand at the top of the image",
|
|
),
|
|
IO.Int.Input(
|
|
"bottom",
|
|
default=0,
|
|
min=0,
|
|
max=2048,
|
|
tooltip="Number of pixels to expand at the bottom of the image",
|
|
),
|
|
IO.Int.Input(
|
|
"left",
|
|
default=0,
|
|
min=0,
|
|
max=2048,
|
|
tooltip="Number of pixels to expand at the left of the image",
|
|
),
|
|
IO.Int.Input(
|
|
"right",
|
|
default=0,
|
|
min=0,
|
|
max=2048,
|
|
tooltip="Number of pixels to expand at the right of the image",
|
|
),
|
|
IO.Float.Input(
|
|
"guidance",
|
|
default=60,
|
|
min=1.5,
|
|
max=100,
|
|
tooltip="Guidance strength for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"steps",
|
|
default=50,
|
|
min=15,
|
|
max=50,
|
|
tooltip="Number of steps for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip="The random seed used for creating the noise.",
|
|
),
|
|
],
|
|
outputs=[IO.Image.Output()],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
image: torch.Tensor,
|
|
prompt: str,
|
|
prompt_upsampling: bool,
|
|
top: int,
|
|
bottom: int,
|
|
left: int,
|
|
right: int,
|
|
steps: int,
|
|
guidance: float,
|
|
seed=0,
|
|
) -> IO.NodeOutput:
|
|
initial_response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/bfl/flux-pro-1.0-expand/generate", method="POST"),
|
|
response_model=BFLFluxProGenerateResponse,
|
|
data=BFLFluxExpandImageRequest(
|
|
prompt=prompt,
|
|
prompt_upsampling=prompt_upsampling,
|
|
top=top,
|
|
bottom=bottom,
|
|
left=left,
|
|
right=right,
|
|
steps=steps,
|
|
guidance=guidance,
|
|
seed=seed,
|
|
image=tensor_to_base64_string(image),
|
|
),
|
|
)
|
|
response = await poll_op(
|
|
cls,
|
|
ApiEndpoint(initial_response.polling_url),
|
|
response_model=BFLFluxStatusResponse,
|
|
status_extractor=lambda r: r.status,
|
|
progress_extractor=lambda r: r.progress,
|
|
completed_statuses=[BFLStatus.ready],
|
|
failed_statuses=[
|
|
BFLStatus.request_moderated,
|
|
BFLStatus.content_moderated,
|
|
BFLStatus.error,
|
|
BFLStatus.task_not_found,
|
|
],
|
|
queued_statuses=[],
|
|
)
|
|
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
|
|
|
|
|
class FluxProFillNode(IO.ComfyNode):
|
|
"""
|
|
Inpaints image based on mask and prompt.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="FluxProFillNode",
|
|
display_name="Flux.1 Fill Image",
|
|
category="api node/image/BFL",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("image"),
|
|
IO.Mask.Input("mask"),
|
|
IO.String.Input(
|
|
"prompt",
|
|
multiline=True,
|
|
default="",
|
|
tooltip="Prompt for the image generation",
|
|
),
|
|
IO.Boolean.Input(
|
|
"prompt_upsampling",
|
|
default=False,
|
|
tooltip="Whether to perform upsampling on the prompt. "
|
|
"If active, automatically modifies the prompt for more creative generation, "
|
|
"but results are nondeterministic (same seed will not produce exactly the same result).",
|
|
),
|
|
IO.Float.Input(
|
|
"guidance",
|
|
default=60,
|
|
min=1.5,
|
|
max=100,
|
|
tooltip="Guidance strength for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"steps",
|
|
default=50,
|
|
min=15,
|
|
max=50,
|
|
tooltip="Number of steps for the image generation process",
|
|
),
|
|
IO.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip="The random seed used for creating the noise.",
|
|
),
|
|
],
|
|
outputs=[IO.Image.Output()],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
image: torch.Tensor,
|
|
mask: torch.Tensor,
|
|
prompt: str,
|
|
prompt_upsampling: bool,
|
|
steps: int,
|
|
guidance: float,
|
|
seed=0,
|
|
) -> IO.NodeOutput:
|
|
# prepare mask
|
|
mask = resize_mask_to_image(mask, image)
|
|
mask = tensor_to_base64_string(convert_mask_to_image(mask))
|
|
initial_response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/bfl/flux-pro-1.0-fill/generate", method="POST"),
|
|
response_model=BFLFluxProGenerateResponse,
|
|
data=BFLFluxFillImageRequest(
|
|
prompt=prompt,
|
|
prompt_upsampling=prompt_upsampling,
|
|
steps=steps,
|
|
guidance=guidance,
|
|
seed=seed,
|
|
image=tensor_to_base64_string(image[:, :, :, :3]), # make sure image will have alpha channel removed
|
|
mask=mask,
|
|
),
|
|
)
|
|
response = await poll_op(
|
|
cls,
|
|
ApiEndpoint(initial_response.polling_url),
|
|
response_model=BFLFluxStatusResponse,
|
|
status_extractor=lambda r: r.status,
|
|
progress_extractor=lambda r: r.progress,
|
|
completed_statuses=[BFLStatus.ready],
|
|
failed_statuses=[
|
|
BFLStatus.request_moderated,
|
|
BFLStatus.content_moderated,
|
|
BFLStatus.error,
|
|
BFLStatus.task_not_found,
|
|
],
|
|
queued_statuses=[],
|
|
)
|
|
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
|
|
|
|
|
class BFLExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
|
return [
|
|
FluxProUltraImageNode,
|
|
# FluxProImageNode,
|
|
FluxKontextProImageNode,
|
|
FluxKontextMaxImageNode,
|
|
FluxProExpandNode,
|
|
FluxProFillNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> BFLExtension:
|
|
return BFLExtension()
|