190 lines
6.5 KiB
Python
190 lines
6.5 KiB
Python
"""Render a preview image of a ComfyUI workflow json.
|
|
|
|
This script *renders* a preview from the workflow graph itself:
|
|
|
|
python tools/render_workflow_preview.py \\
|
|
example_workflows/keep-workflow.json \\
|
|
preview.png
|
|
|
|
The rendered image is a **diagram of the graph** (nodes, titles, widget values
|
|
and the links between them) - not a screenshot of the ComfyUI frontend, and it
|
|
does not pretend to be one. The caption of the image says so.
|
|
|
|
The script only needs Pillow (a runtime dependency of the pack) and the
|
|
standard library.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
BACKGROUND = (24, 24, 32)
|
|
GRID = (34, 34, 44)
|
|
BOX_FILL = (44, 46, 62)
|
|
BOX_OUTLINE = (108, 142, 191)
|
|
BOX_OUTLINE_OUTPUT = (126, 191, 133)
|
|
TITLE_COLOR = (235, 235, 245)
|
|
TYPE_COLOR = (159, 184, 216)
|
|
WIDGET_COLOR = (198, 198, 210)
|
|
LINK_COLOR = (150, 160, 190)
|
|
CAPTION_COLOR = (150, 152, 165)
|
|
|
|
MARGIN = 40
|
|
HEADER = 96
|
|
TARGET_WIDTH = 1600
|
|
|
|
|
|
def _font(size: int):
|
|
"""The scalable default font of Pillow (no font file needed)."""
|
|
try:
|
|
return ImageFont.load_default(size=size)
|
|
except TypeError: # pragma: no cover - Pillow < 10.1
|
|
return ImageFont.load_default()
|
|
|
|
|
|
def _node_title(node: dict) -> str:
|
|
return str(node.get("title") or node.get("type"))
|
|
|
|
|
|
def _fit_text(draw, text: str, font, max_width: float) -> str:
|
|
"""Shorten ``text`` until it fits into ``max_width`` pixels."""
|
|
if draw.textlength(text, font=font) <= max_width:
|
|
return text
|
|
shortened = text
|
|
while shortened and draw.textlength(shortened + "…", font=font) > max_width:
|
|
shortened = shortened[:-1]
|
|
return shortened + "…" if shortened else ""
|
|
|
|
|
|
def _widget_summary(node: dict) -> list[str]:
|
|
"""The widget values of a node as short strings."""
|
|
values = node.get("widgets_values")
|
|
if values is None:
|
|
return []
|
|
if isinstance(values, dict):
|
|
values = list(values.values())
|
|
summary = []
|
|
for value in values:
|
|
if isinstance(value, bool):
|
|
summary.append("true" if value else "false")
|
|
elif value is None:
|
|
summary.append("")
|
|
elif isinstance(value, (int, float)):
|
|
summary.append(str(value))
|
|
else:
|
|
text = str(value)
|
|
summary.append(text if len(text) <= 28 else text[:27] + "…")
|
|
return summary
|
|
|
|
|
|
def _anchor_out(node: dict, index: int) -> tuple[float, float]:
|
|
pos = node.get("pos", [0, 0])
|
|
size = node.get("size", [200, 80])
|
|
return float(pos[0]) + float(size[0]), float(pos[1]) + 26.0 + 18.0 * index
|
|
|
|
|
|
def _anchor_in(node: dict, index: int) -> tuple[float, float]:
|
|
pos = node.get("pos", [0, 0])
|
|
return float(pos[0]), float(pos[1]) + 26.0 + 18.0 * index
|
|
|
|
|
|
def render(workflow_path: Path, output_path: Path) -> Path:
|
|
"""Render the workflow graph to a PNG file."""
|
|
workflow = json.loads(workflow_path.read_text(encoding="utf-8"))
|
|
nodes = workflow.get("nodes", [])
|
|
by_id = {node["id"]: node for node in nodes}
|
|
|
|
max_x = max((float(n.get("pos", [0, 0])[0]) + float(n.get("size", [0, 0])[0]) for n in nodes), default=0.0)
|
|
max_y = max((float(n.get("pos", [0, 0])[1]) + float(n.get("size", [0, 0])[1]) for n in nodes), default=0.0)
|
|
|
|
scale = TARGET_WIDTH / max(max_x + MARGIN * 2, 1.0)
|
|
width = int((max_x + MARGIN * 2) * scale)
|
|
height = int((max_y + MARGIN * 2 + HEADER) * scale)
|
|
|
|
image = Image.new("RGB", (max(width, 640), max(height, 360)), BACKGROUND)
|
|
draw = ImageDraw.Draw(image)
|
|
|
|
title_font = _font(30)
|
|
type_font = _font(19)
|
|
widget_font = _font(16)
|
|
caption_font = _font(15)
|
|
|
|
def point(x: float, y: float) -> tuple[int, int]:
|
|
return int((x + MARGIN) * scale), int((y + MARGIN + HEADER) * scale)
|
|
|
|
# the grid
|
|
step = 40
|
|
x = 0.0
|
|
while x < max_x + MARGIN * 2:
|
|
px = int(x * scale)
|
|
draw.line([(px, 0), (px, image.height)], fill=GRID, width=1)
|
|
x += step
|
|
y = 0.0
|
|
while y < max_y + MARGIN * 2 + HEADER:
|
|
py = int(y * scale)
|
|
draw.line([(0, py), (image.width, py)], fill=GRID, width=1)
|
|
y += step
|
|
|
|
# the links (under the boxes)
|
|
for link in workflow.get("links", []):
|
|
link_id, source_id, source_slot, target_id, target_slot, _kind = link
|
|
source = by_id.get(source_id)
|
|
target = by_id.get(target_id)
|
|
if source is None or target is None:
|
|
continue
|
|
x1, y1 = point(*_anchor_out(source, int(source_slot)))
|
|
x2, y2 = point(*_anchor_in(target, int(target_slot)))
|
|
mid = (x1 + x2) // 2
|
|
draw.line([(x1, y1), (mid, y1), (mid, y2), (x2, y2)], fill=LINK_COLOR, width=2)
|
|
|
|
# the nodes
|
|
for node in nodes:
|
|
pos = node.get("pos", [0, 0])
|
|
size = node.get("size", [200, 80])
|
|
x0, y0 = point(float(pos[0]), float(pos[1]))
|
|
x1, y1 = point(float(pos[0]) + float(size[0]), float(pos[1]) + float(size[1]))
|
|
is_output = node.get("type") in {"SaveImage", "PreviewImage", "VHS_VideoCombine"}
|
|
draw.rounded_rectangle(
|
|
[x0, y0, x1, y1],
|
|
radius=8,
|
|
fill=BOX_FILL,
|
|
outline=BOX_OUTLINE_OUTPUT if is_output else BOX_OUTLINE,
|
|
width=2,
|
|
)
|
|
inner_width = max(x1 - x0 - 24, 20)
|
|
draw.text((x0 + 12, y0 + 8), _fit_text(draw, _node_title(node), title_font, inner_width), font=title_font, fill=TITLE_COLOR)
|
|
draw.text((x0 + 12, y0 + 42), _fit_text(draw, str(node.get("type")), type_font, inner_width), font=type_font, fill=TYPE_COLOR)
|
|
for index, value in enumerate(_widget_summary(node)):
|
|
draw.text((x0 + 12, y0 + 68 + 20 * index), _fit_text(draw, value, widget_font, inner_width), font=widget_font, fill=WIDGET_COLOR)
|
|
|
|
# the header
|
|
draw.text((MARGIN, 18), workflow_path.stem, font=title_font, fill=TITLE_COLOR)
|
|
draw.text(
|
|
(MARGIN, 56),
|
|
"Graph diagram rendered from the workflow json - not a ComfyUI screenshot.",
|
|
font=caption_font,
|
|
fill=CAPTION_COLOR,
|
|
)
|
|
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
image.save(output_path, format="PNG", optimize=True)
|
|
return output_path
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("workflow", type=Path, help="the workflow json to render")
|
|
parser.add_argument("output", type=Path, help="the png to write")
|
|
args = parser.parse_args(argv)
|
|
written = render(args.workflow, args.output)
|
|
print(f"wrote {written}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|