ComfyUI/comfy_extras/nodes_json_prompt.py
Terry Jia 850c91389f
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
feat: Ideogram structured-caption nodes
2026-06-19 22:10:00 -04:00

77 lines
3.5 KiB
Python

from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
from comfy_extras.color_util import normalize_palette
class BuildJsonPromptIdeogram(io.ComfyNode):
@classmethod
def define_schema(cls):
color_palette = io.Colors.Input(
"color_palette",
tooltip="Style color palette.",
)
return io.Schema(
node_id="BuildJsonPromptIdeogram",
display_name="Build JSON Prompt (Ideogram)",
category="image/ideogram",
description="Assemble the Ideogram 4 caption from Create Bounding Boxes elements plus the background and style fields.",
inputs=[
io.ComfyList.Input("element", tooltip="Caption elements from Create Bounding Boxes."),
io.String.Input("high_level_description", multiline=True, default="",
tooltip="Optional one-line overview of the whole image (blank = omitted)."),
io.String.Input("background", multiline=True, default="",
tooltip="Scene background description."),
io.DynamicCombo.Input("style", options=[
io.DynamicCombo.Option("none", []),
io.DynamicCombo.Option("photo", [io.String.Input("photo", default="")]),
io.DynamicCombo.Option("art_style", [io.String.Input("art_style", default="")]),
]),
io.String.Input("aesthetics", default="", tooltip="Style descriptor. Sent even when blank once a style is chosen."),
io.String.Input("lighting", default="", tooltip="Style descriptor. Sent even when blank once a style is chosen."),
io.String.Input("medium", default="", tooltip="Style descriptor. Sent even when blank once a style is chosen."),
color_palette,
],
outputs=[io.ComfyDict.Output(display_name="prompt")],
is_experimental=True,
)
@classmethod
def execute(cls, element, style, high_level_description="", background="",
aesthetics="", lighting="", medium="", color_palette=None) -> io.NodeOutput:
elements = element if isinstance(element, list) else []
kind = style.get("style", "none") if isinstance(style, dict) else "none"
photo = style.get("photo", "") if isinstance(style, dict) else ""
art_style = style.get("art_style", "") if isinstance(style, dict) else ""
palette = normalize_palette(color_palette or [])
caption: dict = {}
if high_level_description.strip():
caption["high_level_description"] = high_level_description
if kind != "none":
style_desc: dict = {"aesthetics": aesthetics, "lighting": lighting}
if kind == "photo":
style_desc["photo"] = photo
style_desc["medium"] = medium
else:
style_desc["medium"] = medium
style_desc["art_style"] = art_style
if palette:
style_desc["color_palette"] = palette
caption["style_description"] = style_desc
caption["compositional_deconstruction"] = {
"background": background,
"elements": elements,
}
return io.NodeOutput(caption)
class JsonPromptExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [BuildJsonPromptIdeogram]
async def comfy_entrypoint() -> JsonPromptExtension:
return JsonPromptExtension()