392 lines
13 KiB
Python
392 lines
13 KiB
Python
"""
|
|
PIL-based grid image assembly with labeled axes.
|
|
|
|
Creates comparison grids where:
|
|
- Rows correspond to Y-axis parameter values
|
|
- Columns correspond to X-axis parameter values
|
|
- Each cell contains the generated image
|
|
- Row and column headers show parameter names and values
|
|
"""
|
|
|
|
import math
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
class GridBuilder:
|
|
def __init__(
|
|
self,
|
|
cell_width: int = 512,
|
|
cell_height: int = 512,
|
|
padding: int = 10,
|
|
label_height: int = 40,
|
|
font_size: int = 14,
|
|
bg_color: tuple = (32, 32, 32),
|
|
text_color: tuple = (255, 255, 255),
|
|
border_color: tuple = (64, 64, 64),
|
|
border_width: int = 1,
|
|
):
|
|
self.cell_width = cell_width
|
|
self.cell_height = cell_height
|
|
self.padding = padding
|
|
self.label_height = label_height
|
|
self.font_size = font_size
|
|
self.bg_color = bg_color
|
|
self.text_color = text_color
|
|
self.border_color = border_color
|
|
self.border_width = border_width
|
|
self._font = None
|
|
|
|
@property
|
|
def font(self) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
|
|
if self._font is None:
|
|
try:
|
|
self._font = ImageFont.truetype("arial.ttf", self.font_size)
|
|
except (OSError, IOError):
|
|
try:
|
|
self._font = ImageFont.truetype(
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
|
self.font_size,
|
|
)
|
|
except (OSError, IOError):
|
|
self._font = ImageFont.load_default()
|
|
return self._font
|
|
|
|
def build_2d_grid(
|
|
self,
|
|
images: list,
|
|
x_labels: list[str],
|
|
y_labels: list[str],
|
|
x_param_name: str = "",
|
|
y_param_name: str = "",
|
|
) -> Image.Image:
|
|
"""
|
|
Build a 2D labeled grid.
|
|
|
|
Args:
|
|
images: List of PIL Images, ordered row-major.
|
|
images[row * num_cols + col] = image at (row, col)
|
|
x_labels: Labels for column headers
|
|
y_labels: Labels for row headers
|
|
x_param_name: Name above column headers (e.g., "Steps")
|
|
y_param_name: Name beside row headers (e.g., "CFG")
|
|
|
|
Returns:
|
|
Single PIL Image with the complete grid.
|
|
"""
|
|
num_cols = len(x_labels)
|
|
num_rows = len(y_labels)
|
|
|
|
# Calculate row label width
|
|
row_label_width = self._calculate_label_width(
|
|
y_labels, y_param_name
|
|
)
|
|
|
|
# Calculate column header height
|
|
col_header_height = self.label_height
|
|
if x_param_name:
|
|
col_header_height += self.label_height
|
|
|
|
# Total dimensions
|
|
total_width = (
|
|
row_label_width
|
|
+ (self.cell_width + self.padding) * num_cols
|
|
+ self.padding
|
|
)
|
|
total_height = (
|
|
col_header_height
|
|
+ (self.cell_height + self.padding) * num_rows
|
|
+ self.padding
|
|
)
|
|
|
|
canvas = Image.new("RGB", (total_width, total_height), self.bg_color)
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
# Draw parameter name headers
|
|
if x_param_name:
|
|
x_center = row_label_width + (
|
|
(self.cell_width + self.padding) * num_cols
|
|
) // 2
|
|
self._draw_centered_text(
|
|
draw, x_param_name, x_center, self.padding, total_width
|
|
)
|
|
|
|
if y_param_name:
|
|
y_center = col_header_height + (
|
|
(self.cell_height + self.padding) * num_rows
|
|
) // 2
|
|
self._draw_rotated_text(
|
|
canvas, y_param_name, self.padding, y_center
|
|
)
|
|
|
|
# Draw column headers (x-axis labels)
|
|
header_y = self.label_height if x_param_name else 0
|
|
for col_idx, label in enumerate(x_labels):
|
|
x = (
|
|
row_label_width
|
|
+ self.padding
|
|
+ col_idx * (self.cell_width + self.padding)
|
|
+ self.cell_width // 2
|
|
)
|
|
self._draw_centered_text(
|
|
draw, str(label), x, header_y + self.padding, self.cell_width
|
|
)
|
|
|
|
# Draw row headers (y-axis labels) and images
|
|
for row_idx, y_label in enumerate(y_labels):
|
|
# Row label
|
|
y = (
|
|
col_header_height
|
|
+ self.padding
|
|
+ row_idx * (self.cell_height + self.padding)
|
|
+ self.cell_height // 2
|
|
)
|
|
label_x = self.padding + (20 if y_param_name else 0)
|
|
self._draw_label(
|
|
draw,
|
|
str(y_label),
|
|
(label_x, y - self.font_size // 2),
|
|
row_label_width - self.padding * 2,
|
|
)
|
|
|
|
for col_idx in range(num_cols):
|
|
img_idx = row_idx * num_cols + col_idx
|
|
if img_idx < len(images) and images[img_idx] is not None:
|
|
cell_img = self._resize_image(
|
|
images[img_idx], self.cell_width, self.cell_height
|
|
)
|
|
|
|
paste_x = (
|
|
row_label_width
|
|
+ self.padding
|
|
+ col_idx * (self.cell_width + self.padding)
|
|
)
|
|
paste_y = (
|
|
col_header_height
|
|
+ self.padding
|
|
+ row_idx * (self.cell_height + self.padding)
|
|
)
|
|
|
|
# Draw border
|
|
if self.border_width > 0:
|
|
bw = self.border_width
|
|
draw.rectangle(
|
|
[
|
|
paste_x - bw,
|
|
paste_y - bw,
|
|
paste_x + self.cell_width + bw - 1,
|
|
paste_y + self.cell_height + bw - 1,
|
|
],
|
|
outline=self.border_color,
|
|
width=bw,
|
|
)
|
|
|
|
canvas.paste(cell_img, (paste_x, paste_y))
|
|
else:
|
|
# Draw placeholder for missing image
|
|
paste_x = (
|
|
row_label_width
|
|
+ self.padding
|
|
+ col_idx * (self.cell_width + self.padding)
|
|
)
|
|
paste_y = (
|
|
col_header_height
|
|
+ self.padding
|
|
+ row_idx * (self.cell_height + self.padding)
|
|
)
|
|
draw.rectangle(
|
|
[
|
|
paste_x,
|
|
paste_y,
|
|
paste_x + self.cell_width - 1,
|
|
paste_y + self.cell_height - 1,
|
|
],
|
|
fill=(64, 0, 0),
|
|
outline=self.border_color,
|
|
)
|
|
self._draw_centered_text(
|
|
draw,
|
|
"ERROR",
|
|
paste_x + self.cell_width // 2,
|
|
paste_y + self.cell_height // 2 - self.font_size // 2,
|
|
self.cell_width,
|
|
)
|
|
|
|
return canvas
|
|
|
|
def build_linear_strip(
|
|
self,
|
|
images: list,
|
|
labels: list[dict],
|
|
columns: int = 0,
|
|
) -> Image.Image:
|
|
"""
|
|
Build a strip/grid from linearly ordered images with per-cell labels.
|
|
|
|
Args:
|
|
images: List of PIL Images
|
|
labels: List of dicts, one per image,
|
|
e.g. [{"steps": 10, "cfg": 5}, ...]
|
|
columns: 0 = auto-calculate, otherwise wrap at this column count
|
|
"""
|
|
count = len(images)
|
|
if count == 0:
|
|
return Image.new("RGB", (100, 100), self.bg_color)
|
|
|
|
if columns <= 0:
|
|
columns = math.ceil(math.sqrt(count))
|
|
|
|
rows = math.ceil(count / columns)
|
|
|
|
# Calculate label text height (multiple params per cell)
|
|
max_params = max(len(lbl) for lbl in labels) if labels else 0
|
|
cell_label_height = max(self.label_height, max_params * (self.font_size + 4) + 8)
|
|
|
|
total_width = (
|
|
self.padding
|
|
+ (self.cell_width + self.padding) * columns
|
|
)
|
|
total_height = (
|
|
self.padding
|
|
+ (self.cell_height + cell_label_height + self.padding) * rows
|
|
)
|
|
|
|
canvas = Image.new("RGB", (total_width, total_height), self.bg_color)
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
for idx in range(count):
|
|
row = idx // columns
|
|
col = idx % columns
|
|
|
|
paste_x = self.padding + col * (self.cell_width + self.padding)
|
|
paste_y = self.padding + row * (
|
|
self.cell_height + cell_label_height + self.padding
|
|
)
|
|
|
|
# Draw image
|
|
if images[idx] is not None:
|
|
cell_img = self._resize_image(
|
|
images[idx], self.cell_width, self.cell_height
|
|
)
|
|
if self.border_width > 0:
|
|
bw = self.border_width
|
|
draw.rectangle(
|
|
[
|
|
paste_x - bw,
|
|
paste_y - bw,
|
|
paste_x + self.cell_width + bw - 1,
|
|
paste_y + self.cell_height + bw - 1,
|
|
],
|
|
outline=self.border_color,
|
|
width=bw,
|
|
)
|
|
canvas.paste(cell_img, (paste_x, paste_y))
|
|
|
|
# Draw labels below image
|
|
if idx < len(labels):
|
|
label_y = paste_y + self.cell_height + 4
|
|
for key, value in labels[idx].items():
|
|
label_text = f"{key}: {value}"
|
|
self._draw_label(
|
|
draw,
|
|
label_text,
|
|
(paste_x + 4, label_y),
|
|
self.cell_width - 8,
|
|
)
|
|
label_y += self.font_size + 4
|
|
|
|
return canvas
|
|
|
|
def _calculate_label_width(
|
|
self, labels: list[str], param_name: str = ""
|
|
) -> int:
|
|
"""Calculate the width needed for row labels."""
|
|
max_text = max(
|
|
(len(str(lbl)) for lbl in labels), default=0
|
|
)
|
|
if param_name:
|
|
max_text = max(max_text, len(param_name))
|
|
# Approximate: font_size * 0.6 per character + padding
|
|
width = int(max_text * self.font_size * 0.65) + self.padding * 3
|
|
return max(width, 60)
|
|
|
|
def _draw_label(
|
|
self, draw: ImageDraw.Draw, text: str, position: tuple, max_width: int
|
|
):
|
|
"""Draw a text label, truncating if necessary."""
|
|
# Truncate text if too long
|
|
while len(text) > 3:
|
|
bbox = self.font.getbbox(text)
|
|
text_width = bbox[2] - bbox[0]
|
|
if text_width <= max_width:
|
|
break
|
|
text = text[:-4] + "..."
|
|
|
|
draw.text(position, text, fill=self.text_color, font=self.font)
|
|
|
|
def _draw_centered_text(
|
|
self, draw: ImageDraw.Draw, text: str, center_x: int, y: int, max_width: int
|
|
):
|
|
"""Draw text centered at center_x."""
|
|
bbox = self.font.getbbox(text)
|
|
text_width = bbox[2] - bbox[0]
|
|
x = center_x - text_width // 2
|
|
draw.text((x, y), text, fill=self.text_color, font=self.font)
|
|
|
|
def _draw_rotated_text(
|
|
self, canvas: Image.Image, text: str, x: int, center_y: int
|
|
):
|
|
"""Draw vertically rotated text."""
|
|
txt_img = Image.new("RGBA", (200, 40), (0, 0, 0, 0))
|
|
txt_draw = ImageDraw.Draw(txt_img)
|
|
txt_draw.text((0, 0), text, fill=self.text_color, font=self.font)
|
|
|
|
bbox = self.font.getbbox(text)
|
|
text_width = bbox[2] - bbox[0]
|
|
text_height = bbox[3] - bbox[1]
|
|
|
|
txt_img = txt_img.crop((0, 0, text_width + 4, text_height + 4))
|
|
txt_img = txt_img.rotate(90, expand=True)
|
|
|
|
paste_y = center_y - txt_img.height // 2
|
|
canvas.paste(txt_img, (x, max(0, paste_y)), txt_img)
|
|
|
|
def _resize_image(
|
|
self, img: Image.Image, target_w: int, target_h: int
|
|
) -> Image.Image:
|
|
"""Resize image to fit cell, maintaining aspect ratio with letterboxing."""
|
|
if img.size == (target_w, target_h):
|
|
return img
|
|
|
|
ratio = min(target_w / img.width, target_h / img.height)
|
|
new_w = int(img.width * ratio)
|
|
new_h = int(img.height * ratio)
|
|
|
|
resized = img.resize((new_w, new_h), Image.LANCZOS)
|
|
|
|
# Letterbox onto background
|
|
result = Image.new("RGB", (target_w, target_h), self.bg_color)
|
|
offset_x = (target_w - new_w) // 2
|
|
offset_y = (target_h - new_h) // 2
|
|
result.paste(resized, (offset_x, offset_y))
|
|
|
|
return result
|
|
|
|
@staticmethod
|
|
def pil_to_tensor(pil_image: Image.Image):
|
|
"""Convert PIL Image to ComfyUI IMAGE tensor (B,H,W,C float32 0-1)."""
|
|
import numpy as np
|
|
import torch
|
|
|
|
img_array = np.array(pil_image.convert("RGB")).astype(np.float32) / 255.0
|
|
return torch.from_numpy(img_array).unsqueeze(0)
|
|
|
|
@staticmethod
|
|
def tensor_to_pil(tensor) -> Image.Image:
|
|
"""Convert ComfyUI IMAGE tensor to PIL Image."""
|
|
import numpy as np
|
|
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
img_array = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
|
return Image.fromarray(img_array)
|