220 lines
7.9 KiB
Python
220 lines
7.9 KiB
Python
"""Grid composition and labeling for PromptMatrix node."""
|
|
|
|
import math
|
|
from typing import Optional
|
|
|
|
_X_LABEL_MAX_LINES = 5
|
|
_LABEL_V_PAD = 10 # px above and below label text in the top strip
|
|
_LINE_SPACING = 3 # px between lines within a label
|
|
_Y_LABEL_H_PAD = 10 # px on each side of text within the left strip
|
|
_Y_STRIP_MIN = 80 # minimum left strip width in px
|
|
_Y_STRIP_MAX = 400 # maximum left strip width in px
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
def _get_font(size: int):
|
|
"""Load a font, falling back to default if needed."""
|
|
try:
|
|
return ImageFont.truetype("DejaVuSans.ttf", size=size)
|
|
except (OSError, IOError):
|
|
# Fall back to default if DejaVuSans not available
|
|
return ImageFont.load_default()
|
|
|
|
|
|
def _wrap_single_line(text: str, max_width: float, font) -> list[str]:
|
|
"""Split a single line of text (no newlines) into word-wrapped segments."""
|
|
lines = []
|
|
current = ""
|
|
for word in text.split():
|
|
test = (current + " " + word).strip() if current else word
|
|
bbox = font.getbbox(test)
|
|
w = bbox[2] - bbox[0]
|
|
if w <= max_width or not current:
|
|
if not current and w > max_width:
|
|
print(
|
|
f"[ComfyUI PromptMatrix] Warning: word '{word}' "
|
|
f"({w}px) exceeds wrap width ({int(max_width)}px) and will overflow"
|
|
)
|
|
current = test
|
|
else:
|
|
lines.append(current)
|
|
current = word
|
|
if current:
|
|
lines.append(current)
|
|
return lines
|
|
|
|
|
|
def _wrap_text(text: str, max_width: float, font, max_lines: int = 2) -> list[str]:
|
|
"""Wrap text to fit within max_width, truncate to max_lines with ellipsis."""
|
|
if not text:
|
|
return []
|
|
|
|
all_wrapped: list[str] = []
|
|
for segment in text.split("\n"):
|
|
all_wrapped.extend(_wrap_single_line(segment, max_width, font))
|
|
|
|
if len(all_wrapped) <= max_lines:
|
|
return all_wrapped
|
|
|
|
lines = all_wrapped[:max_lines]
|
|
last = lines[-1]
|
|
while last and font.getbbox(last + "…")[2] - font.getbbox(last + "…")[0] > max_width:
|
|
last = last[:-1]
|
|
lines[-1] = last + "…"
|
|
return lines
|
|
|
|
|
|
def _measure_text_height(text_lines: list[str], font) -> int:
|
|
"""Measure total height of text lines including inter-line spacing."""
|
|
if not text_lines:
|
|
return 0
|
|
total = sum(font.getbbox(line)[3] - font.getbbox(line)[1] for line in text_lines)
|
|
total += (len(text_lines) - 1) * _LINE_SPACING
|
|
return total
|
|
|
|
|
|
def compose_grid(
|
|
cells: list[torch.Tensor],
|
|
x_labels: Optional[list[str]] = None,
|
|
y_labels: Optional[list[str]] = None,
|
|
width: int = 512,
|
|
height: int = 512,
|
|
max_composite_width: int = 2048,
|
|
) -> torch.Tensor:
|
|
"""Compose individual cell images into a labeled grid.
|
|
|
|
Args:
|
|
cells: list of [H, W, 3] float32 torch tensors (RGB, [0,1])
|
|
x_labels: list of labels for X axis (columns) or None
|
|
y_labels: list of labels for Y axis (rows) or None
|
|
width: original width of each cell before scaling
|
|
height: original height of each cell before scaling
|
|
max_composite_width: maximum width of composite grid
|
|
|
|
Returns:
|
|
[H_grid, W_grid, 3] float32 torch tensor (RGB, [0,1])
|
|
"""
|
|
if not cells:
|
|
raise ValueError("No cells provided")
|
|
|
|
# Determine grid dimensions
|
|
n_x = len(x_labels) if x_labels else 1
|
|
n_y = len(y_labels) if y_labels else 1
|
|
|
|
if len(cells) != n_x * n_y:
|
|
raise ValueError(f"Expected {n_x * n_y} cells, got {len(cells)}")
|
|
|
|
# Calculate cell scale
|
|
n_cells = max(n_x, n_y)
|
|
scale = max(1.0 / math.sqrt(n_cells), 256.0 / max(width, height))
|
|
scale = min(scale, 1.0)
|
|
|
|
cell_w = int(width * scale)
|
|
cell_h = int(height * scale)
|
|
|
|
# Estimate label sizes
|
|
font_size = max(10, min(32, cell_w // 12))
|
|
font = _get_font(font_size)
|
|
|
|
# Calculate label strip dimensions
|
|
if x_labels:
|
|
sample_line_h = font.getbbox("Ag")[3] - font.getbbox("Ag")[1]
|
|
top_strip_height = (
|
|
_X_LABEL_MAX_LINES * sample_line_h
|
|
+ (_X_LABEL_MAX_LINES - 1) * _LINE_SPACING
|
|
+ 2 * _LABEL_V_PAD
|
|
)
|
|
else:
|
|
top_strip_height = 0
|
|
|
|
if y_labels:
|
|
# Measure the single-line pixel width of each label, then clamp.
|
|
# Single-line width is the "ideal" width before any wrapping; clamping
|
|
# gives a strip that fits short labels snugly and caps runaway long ones.
|
|
max_label_px = max(
|
|
font.getbbox(label)[2] - font.getbbox(label)[0] for label in y_labels
|
|
)
|
|
left_strip_width = min(_Y_STRIP_MAX, max(_Y_STRIP_MIN, max_label_px + 2 * _Y_LABEL_H_PAD))
|
|
else:
|
|
left_strip_width = 0
|
|
|
|
# Check if we need to rescale to fit max_composite_width
|
|
tentative_grid_width = left_strip_width + n_x * cell_w
|
|
if tentative_grid_width > max_composite_width and cell_w > 256:
|
|
# Rescale cells to fit
|
|
scale_factor = max_composite_width / tentative_grid_width
|
|
cell_w = int(cell_w * scale_factor)
|
|
cell_h = int(cell_h * scale_factor)
|
|
|
|
if cell_w < 256:
|
|
cell_h = int(cell_h * 256 / cell_w)
|
|
cell_w = 256
|
|
print(f"[ComfyUI PromptMatrix] Warning: resized cells to {cell_w}px, composite may exceed max_composite_width")
|
|
|
|
# Calculate total grid size
|
|
grid_width = left_strip_width + n_x * cell_w
|
|
grid_height = top_strip_height + n_y * cell_h
|
|
|
|
# Create grid image
|
|
grid_img = Image.new("RGB", (grid_width, grid_height), color=(255, 255, 255))
|
|
draw = ImageDraw.Draw(grid_img)
|
|
|
|
# Draw cells and labels
|
|
for yi in range(n_y):
|
|
for xi in range(n_x):
|
|
cell_idx = yi * n_x + xi
|
|
cell_tensor = cells[cell_idx] # [H, W, 3]
|
|
|
|
# Convert tensor to PIL Image
|
|
cell_np = (cell_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
cell_pil = Image.fromarray(cell_np, mode="RGB")
|
|
|
|
# Resize cell
|
|
cell_pil_resized = cell_pil.resize((cell_w, cell_h), Image.Resampling.LANCZOS)
|
|
|
|
# Paste cell into grid
|
|
x_offset = left_strip_width + xi * cell_w
|
|
y_offset = top_strip_height + yi * cell_h
|
|
grid_img.paste(cell_pil_resized, (x_offset, y_offset))
|
|
|
|
# Draw X-axis labels (top strip)
|
|
if x_labels:
|
|
label_y = _LABEL_V_PAD
|
|
for xi, label in enumerate(x_labels):
|
|
x_offset = left_strip_width + xi * cell_w + cell_w // 2
|
|
text_lines = _wrap_text(label, cell_w - 20, font, max_lines=_X_LABEL_MAX_LINES)
|
|
cur_y = label_y
|
|
for line in text_lines:
|
|
bbox = font.getbbox(line)
|
|
line_h = bbox[3] - bbox[1]
|
|
text_x = x_offset - (bbox[2] - bbox[0]) // 2
|
|
draw.text((text_x, cur_y), line, fill=(0, 0, 0), font=font)
|
|
cur_y += line_h + _LINE_SPACING
|
|
|
|
# Draw Y-axis labels (left strip)
|
|
if y_labels:
|
|
wrap_width = left_strip_width - 2 * _Y_LABEL_H_PAD
|
|
for yi, label in enumerate(y_labels):
|
|
y_offset = top_strip_height + yi * cell_h + cell_h // 2
|
|
text_lines = _wrap_text(label, wrap_width, font, max_lines=5)
|
|
|
|
# Center text vertically in the cell; clamp so the first row
|
|
# never renders above the top strip boundary.
|
|
total_height = _measure_text_height(text_lines, font)
|
|
cur_y = max(top_strip_height, y_offset - total_height // 2)
|
|
|
|
for line in text_lines:
|
|
bbox = font.getbbox(line)
|
|
line_h = bbox[3] - bbox[1]
|
|
draw.text((_Y_LABEL_H_PAD, cur_y), line, fill=(0, 0, 0), font=font)
|
|
cur_y += line_h + _LINE_SPACING
|
|
|
|
# Convert grid back to torch tensor
|
|
grid_np = np.array(grid_img, dtype=np.float32) / 255.0
|
|
grid_tensor = torch.from_numpy(grid_np)
|
|
|
|
return grid_tensor
|