1197 lines
39 KiB
Python
1197 lines
39 KiB
Python
import copy
|
|
import gc
|
|
import math
|
|
import textwrap
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
import comfy.model_management
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import folder_paths
|
|
from nodes import CheckpointLoaderSimple, UNETLoader, common_ksampler
|
|
|
|
|
|
MAX_SEED = 0xFFFFFFFFFFFFFFFF
|
|
DEFAULT_MAX_COMBINATIONS = 100
|
|
NONE_OPTION = "None"
|
|
SAMPLER_SLOT_COUNT = 9
|
|
SCHEDULER_SLOT_COUNT = 9
|
|
MODEL_SLOT_COUNT = 20
|
|
UNKNOWN_LABEL = "unknown"
|
|
CHECKPOINT_PREFIX = "checkpoint | "
|
|
DIFFUSION_PREFIX = "diffusion | "
|
|
|
|
|
|
def collect_selected_slots(values):
|
|
seen = set()
|
|
selected = []
|
|
|
|
for value in values:
|
|
name = str(value).strip() if value is not None else NONE_OPTION
|
|
if not name or name == NONE_OPTION or name.lower() == "none" or name in seen:
|
|
continue
|
|
seen.add(name)
|
|
selected.append(name)
|
|
|
|
return selected
|
|
|
|
|
|
def get_available_samplers():
|
|
return list(comfy.samplers.KSampler.SAMPLERS)
|
|
|
|
|
|
def get_available_schedulers():
|
|
return list(comfy.samplers.KSampler.SCHEDULERS)
|
|
|
|
|
|
def get_available_model_choices():
|
|
checkpoints = [f"{CHECKPOINT_PREFIX}{name}" for name in folder_paths.get_filename_list("checkpoints")]
|
|
diffusion_models = [
|
|
f"{DIFFUSION_PREFIX}{name}"
|
|
for name in folder_paths.get_filename_list("diffusion_models")
|
|
]
|
|
return checkpoints + diffusion_models
|
|
|
|
|
|
def parse_model_selection(selection):
|
|
if selection.startswith(CHECKPOINT_PREFIX):
|
|
return "checkpoint", selection[len(CHECKPOINT_PREFIX):]
|
|
if selection.startswith(DIFFUSION_PREFIX):
|
|
return "diffusion", selection[len(DIFFUSION_PREFIX):]
|
|
raise ValueError(f"Unsupported model selection: {selection}")
|
|
|
|
|
|
def load_model_selection(selection, diffusion_weight_dtype):
|
|
source_type, model_name = parse_model_selection(selection)
|
|
if source_type == "checkpoint":
|
|
model, clip, vae = CheckpointLoaderSimple().load_checkpoint(model_name)
|
|
return model, clip, vae, source_type, model_name
|
|
|
|
model = UNETLoader().load_unet(model_name, diffusion_weight_dtype)[0]
|
|
return model, None, None, source_type, model_name
|
|
|
|
|
|
def encode_prompt_text(clip, text):
|
|
if not is_valid_clip(clip):
|
|
raise RuntimeError(
|
|
"No compatible CLIP is available. Connect the optional CLIP input "
|
|
"when comparing standalone diffusion models."
|
|
)
|
|
tokens = clip.tokenize(text)
|
|
return clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
|
|
def is_valid_clip(clip):
|
|
if clip is None:
|
|
return False
|
|
return callable(getattr(clip, "tokenize", None)) and callable(
|
|
getattr(clip, "encode_from_tokens_scheduled", None)
|
|
)
|
|
|
|
|
|
def is_valid_vae(vae):
|
|
if vae is None or not callable(getattr(vae, "decode", None)):
|
|
return False
|
|
|
|
validator = getattr(vae, "throw_exception_if_invalid", None)
|
|
if callable(validator):
|
|
try:
|
|
validator()
|
|
except RuntimeError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def select_clip(checkpoint_clip, external_clip):
|
|
if is_valid_clip(checkpoint_clip):
|
|
return checkpoint_clip
|
|
if is_valid_clip(external_clip):
|
|
return external_clip
|
|
return None
|
|
|
|
|
|
def select_vae(checkpoint_vae, external_vae):
|
|
if is_valid_vae(checkpoint_vae):
|
|
return checkpoint_vae
|
|
if is_valid_vae(external_vae):
|
|
return external_vae
|
|
return None
|
|
|
|
|
|
def validate_selected_names(selected, available, label):
|
|
if not selected:
|
|
raise ValueError(f"No {label} selected. Choose at least one {label} slot that is not None.")
|
|
|
|
available_set = set(available)
|
|
unknown = [name for name in selected if name not in available_set]
|
|
if unknown:
|
|
available_text = ", ".join(available)
|
|
unknown_text = ", ".join(unknown)
|
|
raise ValueError(
|
|
f"Unknown {label}: {unknown_text}. Available {label}: {available_text}"
|
|
)
|
|
|
|
|
|
def get_prompt_node(prompt, node_id):
|
|
if not isinstance(prompt, dict) or node_id is None:
|
|
return None
|
|
|
|
return prompt.get(str(node_id)) or prompt.get(node_id)
|
|
|
|
|
|
def get_input_link(prompt, node_id, input_name):
|
|
node = get_prompt_node(prompt, node_id)
|
|
if not isinstance(node, dict):
|
|
return None
|
|
|
|
value = node.get("inputs", {}).get(input_name)
|
|
if isinstance(value, (list, tuple)) and value:
|
|
return value[0]
|
|
return None
|
|
|
|
|
|
def get_first_widget_value(node, names):
|
|
if not isinstance(node, dict):
|
|
return None
|
|
|
|
inputs = node.get("inputs", {})
|
|
for name in names:
|
|
value = inputs.get(name)
|
|
if isinstance(value, (str, int, float)) and str(value).strip():
|
|
return str(value).strip()
|
|
return None
|
|
|
|
|
|
def describe_source_node(node, preferred_inputs):
|
|
value = get_first_widget_value(node, preferred_inputs)
|
|
if value:
|
|
return value
|
|
|
|
if isinstance(node, dict):
|
|
class_type = node.get("class_type")
|
|
if class_type:
|
|
return str(class_type)
|
|
|
|
return UNKNOWN_LABEL
|
|
|
|
|
|
def infer_clip_label(prompt, unique_id, model_label):
|
|
for conditioning_input in ("positive", "negative"):
|
|
conditioning_node_id = get_input_link(prompt, unique_id, conditioning_input)
|
|
clip_source_id = get_input_link(prompt, conditioning_node_id, "clip")
|
|
clip_source = get_prompt_node(prompt, clip_source_id)
|
|
clip_label = describe_source_node(
|
|
clip_source,
|
|
("clip_name", "clip_name1", "clip_name2", "ckpt_name", "checkpoint", "model_name"),
|
|
)
|
|
if clip_label != UNKNOWN_LABEL:
|
|
if model_label != UNKNOWN_LABEL and clip_label == model_label:
|
|
return "same as model checkpoint"
|
|
return clip_label
|
|
|
|
return UNKNOWN_LABEL
|
|
|
|
|
|
def infer_run_metadata(prompt, unique_id, steps, cfg, denoise):
|
|
model_source_id = get_input_link(prompt, unique_id, "model")
|
|
vae_source_id = get_input_link(prompt, unique_id, "vae")
|
|
model_source = get_prompt_node(prompt, model_source_id)
|
|
vae_source = get_prompt_node(prompt, vae_source_id)
|
|
|
|
model_label = describe_source_node(
|
|
model_source,
|
|
("ckpt_name", "checkpoint", "unet_name", "model_name", "model"),
|
|
)
|
|
vae_label = describe_source_node(
|
|
vae_source,
|
|
("vae_name", "ckpt_name", "checkpoint", "model_name", "vae"),
|
|
)
|
|
if model_label != UNKNOWN_LABEL and vae_label == model_label:
|
|
vae_label = "same as model checkpoint"
|
|
|
|
clip_label = infer_clip_label(prompt, unique_id, model_label)
|
|
|
|
return {
|
|
"model": model_label,
|
|
"vae": vae_label,
|
|
"clip": clip_label,
|
|
"steps": str(steps),
|
|
"cfg": f"{float(cfg):g}",
|
|
"denoise": f"{float(denoise):g}",
|
|
}
|
|
|
|
|
|
def clone_latent(latent):
|
|
cloned = {}
|
|
for key, value in latent.items():
|
|
if torch.is_tensor(value):
|
|
cloned[key] = value.clone()
|
|
else:
|
|
try:
|
|
cloned[key] = copy.deepcopy(value)
|
|
except Exception:
|
|
cloned[key] = value
|
|
return cloned
|
|
|
|
|
|
def run_single_sample(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise,
|
|
):
|
|
return common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
clone_latent(latent_image),
|
|
denoise=denoise,
|
|
)[0]
|
|
|
|
|
|
def decode_latent_to_image(vae, sampled_latent):
|
|
latent = sampled_latent["samples"]
|
|
if getattr(latent, "is_nested", False):
|
|
latent = latent.unbind()[0]
|
|
|
|
images = vae.decode(latent)
|
|
if len(images.shape) == 5:
|
|
images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
|
|
return images
|
|
|
|
|
|
def tensor_to_pil(image_tensor):
|
|
image = image_tensor
|
|
if len(image.shape) == 4:
|
|
image = image[0]
|
|
|
|
array = (
|
|
image.detach()
|
|
.cpu()
|
|
.clamp(0.0, 1.0)
|
|
.numpy()
|
|
)
|
|
array = (array * 255.0).round().astype(np.uint8)
|
|
|
|
if array.ndim == 2:
|
|
return Image.fromarray(array, mode="L").convert("RGB")
|
|
|
|
if array.shape[-1] == 1:
|
|
return Image.fromarray(array[..., 0], mode="L").convert("RGB")
|
|
|
|
return Image.fromarray(array[..., :3], mode="RGB")
|
|
|
|
|
|
def pil_to_tensor(image):
|
|
image = image.convert("RGB")
|
|
array = np.asarray(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(array).unsqueeze(0)
|
|
|
|
|
|
def resolve_background(background):
|
|
colors = {
|
|
"white": (246, 246, 242),
|
|
"gray": (42, 44, 48),
|
|
"black": (12, 13, 15),
|
|
}
|
|
return colors.get(background, colors["white"])
|
|
|
|
|
|
def text_color_for_background(background):
|
|
if background in ("black", "gray"):
|
|
return (238, 238, 232)
|
|
return (28, 30, 34)
|
|
|
|
|
|
def line_color_for_background(background):
|
|
if background in ("black", "gray"):
|
|
return (92, 96, 104)
|
|
return (184, 184, 176)
|
|
|
|
|
|
def label_background_for_background(background):
|
|
colors = {
|
|
"white": (235, 235, 229),
|
|
"gray": (34, 36, 40),
|
|
"black": (24, 25, 29),
|
|
}
|
|
return colors.get(background, colors["white"])
|
|
|
|
|
|
def load_font(size):
|
|
for font_name in ("arial.ttf", "DejaVuSans.ttf", "LiberationSans-Regular.ttf"):
|
|
try:
|
|
return ImageFont.truetype(font_name, size=size)
|
|
except OSError:
|
|
continue
|
|
return ImageFont.load_default()
|
|
|
|
|
|
def text_size(draw, text, font):
|
|
bbox = draw.textbbox((0, 0), text, font=font)
|
|
return bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|
|
|
|
|
def wrap_text(draw, text, font, max_width):
|
|
if not text:
|
|
return [""]
|
|
|
|
wrapped = []
|
|
for raw_line in str(text).splitlines():
|
|
words = raw_line.split()
|
|
if not words:
|
|
wrapped.append("")
|
|
continue
|
|
|
|
line = words[0]
|
|
for word in words[1:]:
|
|
candidate = f"{line} {word}"
|
|
if text_size(draw, candidate, font)[0] <= max_width:
|
|
line = candidate
|
|
else:
|
|
wrapped.append(line)
|
|
line = word
|
|
|
|
while text_size(draw, line, font)[0] > max_width and len(line) > 1:
|
|
approx_chars = max(1, int(len(line) * max_width / max(text_size(draw, line, font)[0], 1)))
|
|
wrapped.extend(textwrap.wrap(line, width=approx_chars, break_long_words=True))
|
|
line = ""
|
|
|
|
if line:
|
|
wrapped.append(line)
|
|
|
|
return wrapped or [str(text)]
|
|
|
|
|
|
def draw_centered_text(draw, box, text, font, fill):
|
|
x, y, width, height = box
|
|
max_width = max(1, width - 12)
|
|
lines = wrap_text(draw, text, font, max_width)
|
|
line_heights = [text_size(draw, line, font)[1] for line in lines]
|
|
total_height = sum(line_heights) + max(0, len(lines) - 1) * 4
|
|
cursor_y = y + max(0, (height - total_height) // 2)
|
|
|
|
for line, line_height in zip(lines, line_heights):
|
|
line_width, _ = text_size(draw, line, font)
|
|
draw.text((x + (width - line_width) // 2, cursor_y), line, font=font, fill=fill)
|
|
cursor_y += line_height + 4
|
|
|
|
|
|
def draw_left_label(draw, box, text, font, fill):
|
|
x, y, width, height = box
|
|
max_width = max(1, width - 18)
|
|
lines = wrap_text(draw, text, font, max_width)
|
|
line_heights = [text_size(draw, line, font)[1] for line in lines]
|
|
total_height = sum(line_heights) + max(0, len(lines) - 1) * 4
|
|
cursor_y = y + max(0, (height - total_height) // 2)
|
|
|
|
for line, line_height in zip(lines, line_heights):
|
|
draw.text((x + 10, cursor_y), line, font=font, fill=fill)
|
|
cursor_y += line_height + 4
|
|
|
|
|
|
def build_run_header_text(metadata):
|
|
return (
|
|
"KSampler Matrix Lab Benchmark\n"
|
|
f"Model: {metadata['model']} | VAE: {metadata['vae']} | CLIP: {metadata['clip']}\n"
|
|
f"Steps: {metadata['steps']} | CFG: {metadata['cfg']} | Denoise: {metadata['denoise']}"
|
|
)
|
|
|
|
|
|
def wrapped_text_height(draw, text, font, max_width):
|
|
lines = wrap_text(draw, text, font, max_width)
|
|
line_heights = [text_size(draw, line, font)[1] for line in lines]
|
|
return sum(line_heights) + max(0, len(lines) - 1) * 4
|
|
|
|
|
|
def draw_run_header(draw, box, metadata, font, fill):
|
|
text = build_run_header_text(metadata)
|
|
draw_centered_text(draw, box, text, font, fill)
|
|
|
|
|
|
def placeholder_size_from_latent(latent_image):
|
|
samples = latent_image.get("samples")
|
|
if torch.is_tensor(samples) and samples.ndim >= 4:
|
|
return max(1, int(samples.shape[-1]) * 8), max(1, int(samples.shape[-2]) * 8)
|
|
return 512, 512
|
|
|
|
|
|
def compose_error_cell(size, sampler, scheduler, error_message, background, font):
|
|
cell = Image.new("RGB", size, resolve_background(background))
|
|
draw = ImageDraw.Draw(cell)
|
|
fill = (210, 76, 68) if background == "white" else (255, 146, 132)
|
|
short_error = str(error_message).replace("\n", " ")[:140]
|
|
text = f"ERROR\n{sampler} / {scheduler}\n{short_error}"
|
|
draw_centered_text(draw, (8, 8, size[0] - 16, size[1] - 16), text, font, fill)
|
|
return cell
|
|
|
|
|
|
def compose_grid(
|
|
rows,
|
|
columns,
|
|
cells,
|
|
latent_image,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
header_height,
|
|
left_header_width,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
run_metadata,
|
|
):
|
|
font = load_font(font_size)
|
|
label_font = load_font(max(10, int(font_size * 0.92)))
|
|
cell_label_font = load_font(max(9, int(font_size * 0.78)))
|
|
|
|
first_image = next((cell["image"] for cell in cells.values() if cell.get("image") is not None), None)
|
|
base_width, base_height = first_image.size if first_image is not None else placeholder_size_from_latent(latent_image)
|
|
|
|
cell_width = max(1, int(round(base_width * cell_scale)))
|
|
cell_height = max(1, int(round(base_height * cell_scale)))
|
|
cell_label_height = max(42, int(font_size * 2.7)) if show_cell_labels else 0
|
|
cell_block_height = cell_height + cell_label_height
|
|
|
|
background_rgb = resolve_background(background)
|
|
label_background_rgb = label_background_for_background(background)
|
|
text_rgb = text_color_for_background(background)
|
|
line_rgb = line_color_for_background(background)
|
|
|
|
total_width = left_header_width + len(columns) * cell_width + (len(columns) + 1) * padding
|
|
if show_run_header:
|
|
measure = ImageDraw.Draw(Image.new("RGB", (1, 1)))
|
|
header_text_height = wrapped_text_height(
|
|
measure,
|
|
build_run_header_text(run_metadata),
|
|
label_font,
|
|
max(1, total_width - 2 * padding),
|
|
)
|
|
run_header_height = max(70, header_text_height + 2 * padding)
|
|
else:
|
|
run_header_height = 0
|
|
|
|
total_height = run_header_height + header_height + len(rows) * cell_block_height + (len(rows) + 1) * padding
|
|
|
|
grid = Image.new("RGB", (total_width, total_height), background_rgb)
|
|
draw = ImageDraw.Draw(grid)
|
|
|
|
grid_y = run_header_height
|
|
if show_run_header:
|
|
draw.rectangle((0, 0, total_width - 1, run_header_height - 1), fill=label_background_rgb)
|
|
draw_run_header(
|
|
draw,
|
|
(padding, padding, total_width - 2 * padding, max(1, run_header_height - 2 * padding)),
|
|
run_metadata,
|
|
label_font,
|
|
text_rgb,
|
|
)
|
|
if show_grid_lines:
|
|
draw.line((0, run_header_height - 1, total_width, run_header_height - 1), fill=line_rgb, width=1)
|
|
|
|
draw_centered_text(
|
|
draw,
|
|
(0, grid_y, left_header_width, header_height),
|
|
"KSampler\nMatrix Lab",
|
|
label_font,
|
|
text_rgb,
|
|
)
|
|
|
|
for column_index, scheduler in enumerate(columns):
|
|
x = left_header_width + padding + column_index * (cell_width + padding)
|
|
draw_centered_text(draw, (x, grid_y, cell_width, header_height), scheduler, font, text_rgb)
|
|
|
|
for row_index, sampler in enumerate(rows):
|
|
y = grid_y + header_height + padding + row_index * (cell_block_height + padding)
|
|
draw_left_label(draw, (0, y, left_header_width, cell_block_height), sampler, font, text_rgb)
|
|
|
|
for row_index, sampler in enumerate(rows):
|
|
for column_index, scheduler in enumerate(columns):
|
|
x = left_header_width + padding + column_index * (cell_width + padding)
|
|
y = grid_y + header_height + padding + row_index * (cell_block_height + padding)
|
|
cell = cells[(sampler, scheduler)]
|
|
|
|
if cell.get("image") is not None:
|
|
image = cell["image"]
|
|
if image.size != (cell_width, cell_height):
|
|
image = image.resize((cell_width, cell_height), Image.Resampling.LANCZOS)
|
|
else:
|
|
image = compose_error_cell(
|
|
(cell_width, cell_height),
|
|
sampler,
|
|
scheduler,
|
|
cell.get("error", "Unknown error"),
|
|
background,
|
|
label_font,
|
|
)
|
|
|
|
if show_cell_labels:
|
|
draw.rectangle(
|
|
(x, y, x + cell_width - 1, y + cell_label_height - 1),
|
|
fill=label_background_rgb,
|
|
)
|
|
draw_centered_text(
|
|
draw,
|
|
(x + 4, y + 3, cell_width - 8, cell_label_height - 6),
|
|
f"Sampler: {sampler}\nScheduler: {scheduler}",
|
|
cell_label_font,
|
|
text_rgb,
|
|
)
|
|
if show_grid_lines:
|
|
draw.line(
|
|
(x, y + cell_label_height - 1, x + cell_width - 1, y + cell_label_height - 1),
|
|
fill=line_rgb,
|
|
width=1,
|
|
)
|
|
|
|
grid.paste(image, (x, y + cell_label_height))
|
|
|
|
if show_grid_lines:
|
|
draw.rectangle(
|
|
(x, y, x + cell_width - 1, y + cell_block_height - 1),
|
|
outline=line_rgb,
|
|
width=1,
|
|
)
|
|
|
|
if show_grid_lines:
|
|
draw.line((left_header_width, grid_y, left_header_width, total_height), fill=line_rgb, width=1)
|
|
draw.line((0, grid_y + header_height, total_width, grid_y + header_height), fill=line_rgb, width=1)
|
|
|
|
return grid
|
|
|
|
|
|
def build_model_run_header_text(sampler_name, scheduler, steps, cfg, denoise):
|
|
return (
|
|
"Model Matrix Lab Benchmark\n"
|
|
f"Sampler: {sampler_name} | Scheduler: {scheduler}\n"
|
|
f"Steps: {steps} | CFG: {float(cfg):g} | Denoise: {float(denoise):g}"
|
|
)
|
|
|
|
|
|
def compose_model_error_cell(size, model_label, error_message, background, font):
|
|
cell = Image.new("RGB", size, resolve_background(background))
|
|
draw = ImageDraw.Draw(cell)
|
|
fill = (210, 76, 68) if background == "white" else (255, 146, 132)
|
|
short_error = str(error_message).replace("\n", " ")[:180]
|
|
draw_centered_text(
|
|
draw,
|
|
(8, 8, size[0] - 16, size[1] - 16),
|
|
f"ERROR\n{model_label}\n{short_error}",
|
|
font,
|
|
fill,
|
|
)
|
|
return cell
|
|
|
|
|
|
def compose_model_grid(
|
|
entries,
|
|
latent_image,
|
|
sampler_name,
|
|
scheduler,
|
|
steps,
|
|
cfg,
|
|
denoise,
|
|
grid_columns,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
):
|
|
font = load_font(font_size)
|
|
label_font = load_font(max(9, int(font_size * 0.82)))
|
|
first_image = next(
|
|
(entry["image"] for entry in entries if entry.get("image") is not None),
|
|
None,
|
|
)
|
|
base_width, base_height = (
|
|
first_image.size
|
|
if first_image is not None
|
|
else placeholder_size_from_latent(latent_image)
|
|
)
|
|
|
|
cell_width = max(1, int(round(base_width * cell_scale)))
|
|
cell_height = max(1, int(round(base_height * cell_scale)))
|
|
cell_label_height = max(42, int(font_size * 2.7)) if show_cell_labels else 0
|
|
cell_block_height = cell_height + cell_label_height
|
|
|
|
columns = max(1, min(int(grid_columns), len(entries)))
|
|
rows = math.ceil(len(entries) / columns)
|
|
background_rgb = resolve_background(background)
|
|
label_background_rgb = label_background_for_background(background)
|
|
text_rgb = text_color_for_background(background)
|
|
line_rgb = line_color_for_background(background)
|
|
|
|
total_width = columns * cell_width + (columns + 1) * padding
|
|
header_text = build_model_run_header_text(
|
|
sampler_name,
|
|
scheduler,
|
|
steps,
|
|
cfg,
|
|
denoise,
|
|
)
|
|
if show_run_header:
|
|
measure = ImageDraw.Draw(Image.new("RGB", (1, 1)))
|
|
header_text_height = wrapped_text_height(
|
|
measure,
|
|
header_text,
|
|
font,
|
|
max(1, total_width - 2 * padding),
|
|
)
|
|
run_header_height = max(70, header_text_height + 2 * padding)
|
|
else:
|
|
run_header_height = 0
|
|
|
|
total_height = run_header_height + rows * cell_block_height + (rows + 1) * padding
|
|
grid = Image.new("RGB", (total_width, total_height), background_rgb)
|
|
draw = ImageDraw.Draw(grid)
|
|
|
|
if show_run_header:
|
|
draw.rectangle(
|
|
(0, 0, total_width - 1, run_header_height - 1),
|
|
fill=label_background_rgb,
|
|
)
|
|
draw_centered_text(
|
|
draw,
|
|
(padding, padding, total_width - 2 * padding, run_header_height - 2 * padding),
|
|
header_text,
|
|
font,
|
|
text_rgb,
|
|
)
|
|
if show_grid_lines:
|
|
draw.line(
|
|
(0, run_header_height - 1, total_width, run_header_height - 1),
|
|
fill=line_rgb,
|
|
width=1,
|
|
)
|
|
|
|
for index, entry in enumerate(entries):
|
|
row = index // columns
|
|
column = index % columns
|
|
x = padding + column * (cell_width + padding)
|
|
y = run_header_height + padding + row * (cell_block_height + padding)
|
|
|
|
if show_cell_labels:
|
|
draw.rectangle(
|
|
(x, y, x + cell_width - 1, y + cell_label_height - 1),
|
|
fill=label_background_rgb,
|
|
)
|
|
draw_centered_text(
|
|
draw,
|
|
(x + 4, y + 3, cell_width - 8, cell_label_height - 6),
|
|
f"Model: {entry['label']}",
|
|
label_font,
|
|
text_rgb,
|
|
)
|
|
if show_grid_lines:
|
|
draw.line(
|
|
(x, y + cell_label_height - 1, x + cell_width - 1, y + cell_label_height - 1),
|
|
fill=line_rgb,
|
|
width=1,
|
|
)
|
|
|
|
if entry.get("image") is not None:
|
|
image = entry["image"]
|
|
if image.size != (cell_width, cell_height):
|
|
image = image.resize((cell_width, cell_height), Image.Resampling.LANCZOS)
|
|
else:
|
|
image = compose_model_error_cell(
|
|
(cell_width, cell_height),
|
|
entry["label"],
|
|
entry.get("error", "Unknown error"),
|
|
background,
|
|
label_font,
|
|
)
|
|
|
|
grid.paste(image, (x, y + cell_label_height))
|
|
if show_grid_lines:
|
|
draw.rectangle(
|
|
(x, y, x + cell_width - 1, y + cell_block_height - 1),
|
|
outline=line_rgb,
|
|
width=1,
|
|
)
|
|
|
|
return grid
|
|
|
|
|
|
class KSamplerMatrixLab:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
samplers = get_available_samplers()
|
|
schedulers = get_available_schedulers()
|
|
sampler_choices = [NONE_OPTION] + samplers
|
|
scheduler_choices = [NONE_OPTION] + schedulers
|
|
default_sampler = samplers[0] if samplers else NONE_OPTION
|
|
default_scheduler = schedulers[0] if schedulers else NONE_OPTION
|
|
|
|
required = {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_SEED,
|
|
"control_after_generate": True,
|
|
},
|
|
),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
|
|
for index in range(1, SAMPLER_SLOT_COUNT + 1):
|
|
required[f"sampler_{index:02d}"] = (
|
|
sampler_choices,
|
|
{
|
|
"default": default_sampler if index == 1 else NONE_OPTION,
|
|
"tooltip": "Select a sampler for a matrix row, or None to ignore this slot.",
|
|
},
|
|
)
|
|
|
|
for index in range(1, SCHEDULER_SLOT_COUNT + 1):
|
|
required[f"scheduler_{index:02d}"] = (
|
|
scheduler_choices,
|
|
{
|
|
"default": default_scheduler if index == 1 else NONE_OPTION,
|
|
"tooltip": "Select a scheduler for a matrix column, or None to ignore this slot.",
|
|
},
|
|
)
|
|
|
|
required.update(
|
|
{
|
|
"seed_mode": (["same_seed_for_all", "increment_per_cell"], {"default": "same_seed_for_all"}),
|
|
"cell_scale": ("FLOAT", {"default": 1.0, "min": 0.05, "max": 4.0, "step": 0.05}),
|
|
"font_size": ("INT", {"default": 22, "min": 8, "max": 96}),
|
|
"padding": ("INT", {"default": 10, "min": 0, "max": 96}),
|
|
"header_height": ("INT", {"default": 72, "min": 24, "max": 256}),
|
|
"left_header_width": ("INT", {"default": 180, "min": 48, "max": 512}),
|
|
"background": (["white", "gray", "black"], {"default": "white"}),
|
|
"show_grid_lines": ("BOOLEAN", {"default": True}),
|
|
"show_cell_labels": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Repeat sampler and scheduler labels inside every grid cell for easier zoomed inspection.",
|
|
},
|
|
),
|
|
"show_run_header": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Add a top header with model, VAE, CLIP, steps, CFG, and denoise metadata.",
|
|
},
|
|
),
|
|
"continue_on_error": ("BOOLEAN", {"default": True}),
|
|
"max_combinations": (
|
|
"INT",
|
|
{
|
|
"default": DEFAULT_MAX_COMBINATIONS,
|
|
"min": 1,
|
|
"max": 1000,
|
|
"tooltip": "Safety limit for sampler x scheduler combinations.",
|
|
},
|
|
),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"required": required,
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("grid_image",)
|
|
FUNCTION = "generate_matrix"
|
|
CATEGORY = "ComfyUI-KSampler-Matrix-Lab"
|
|
DESCRIPTION = "Runs selected sampler and scheduler combinations sequentially, then returns one labeled benchmark grid image."
|
|
|
|
def generate_matrix(
|
|
self,
|
|
model,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
vae,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
denoise,
|
|
seed_mode,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
header_height,
|
|
left_header_width,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
continue_on_error,
|
|
max_combinations,
|
|
prompt=None,
|
|
unique_id=None,
|
|
**selection_slots,
|
|
):
|
|
selected_samplers = collect_selected_slots(
|
|
selection_slots.get(f"sampler_{index:02d}")
|
|
for index in range(1, SAMPLER_SLOT_COUNT + 1)
|
|
)
|
|
selected_schedulers = collect_selected_slots(
|
|
selection_slots.get(f"scheduler_{index:02d}")
|
|
for index in range(1, SCHEDULER_SLOT_COUNT + 1)
|
|
)
|
|
|
|
validate_selected_names(selected_samplers, get_available_samplers(), "sampler")
|
|
validate_selected_names(selected_schedulers, get_available_schedulers(), "scheduler")
|
|
|
|
total = len(selected_samplers) * len(selected_schedulers)
|
|
if total > max_combinations:
|
|
raise ValueError(
|
|
f"KSampler Matrix Lab refused to run {total} combinations. "
|
|
f"Current max_combinations is {max_combinations}."
|
|
)
|
|
|
|
progress = comfy.utils.ProgressBar(total)
|
|
cells = {}
|
|
cell_index = 0
|
|
|
|
for sampler_name in selected_samplers:
|
|
for scheduler in selected_schedulers:
|
|
current_seed = seed if seed_mode == "same_seed_for_all" else (seed + cell_index) & MAX_SEED
|
|
try:
|
|
sampled_latent = run_single_sample(
|
|
model,
|
|
current_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise,
|
|
)
|
|
decoded_images = decode_latent_to_image(vae, sampled_latent)
|
|
cells[(sampler_name, scheduler)] = {"image": tensor_to_pil(decoded_images)}
|
|
except Exception as exc:
|
|
if not continue_on_error:
|
|
raise RuntimeError(
|
|
f"KSampler Matrix Lab failed for sampler '{sampler_name}' "
|
|
f"and scheduler '{scheduler}': {exc}"
|
|
) from exc
|
|
cells[(sampler_name, scheduler)] = {"image": None, "error": str(exc)}
|
|
finally:
|
|
cell_index += 1
|
|
progress.update(1)
|
|
gc.collect()
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
grid = compose_grid(
|
|
selected_samplers,
|
|
selected_schedulers,
|
|
cells,
|
|
latent_image,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
header_height,
|
|
left_header_width,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
infer_run_metadata(prompt, unique_id, steps, cfg, denoise),
|
|
)
|
|
|
|
return (pil_to_tensor(grid),)
|
|
|
|
|
|
class ModelMatrixLab:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
model_choices = [NONE_OPTION] + get_available_model_choices()
|
|
available_models = get_available_model_choices()
|
|
samplers = get_available_samplers()
|
|
schedulers = get_available_schedulers()
|
|
|
|
required = {
|
|
"latent_image": ("LATENT",),
|
|
"positive_text": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"dynamicPrompts": True,
|
|
"tooltip": "The same positive prompt is encoded for every selected model.",
|
|
},
|
|
),
|
|
"negative_text": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"dynamicPrompts": True,
|
|
"tooltip": "The same negative prompt is encoded for every selected model.",
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": MAX_SEED,
|
|
"control_after_generate": True,
|
|
},
|
|
),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"sampler_name": (
|
|
samplers,
|
|
{"tooltip": "The sampler used for every selected model."},
|
|
),
|
|
"scheduler": (
|
|
schedulers,
|
|
{"tooltip": "The scheduler used for every selected model."},
|
|
),
|
|
}
|
|
|
|
for index in range(1, MODEL_SLOT_COUNT + 1):
|
|
required[f"model_{index:02d}"] = (
|
|
model_choices,
|
|
{
|
|
"default": available_models[0] if index == 1 and available_models else NONE_OPTION,
|
|
"tooltip": "Select a checkpoint or standalone diffusion model, or None to ignore this slot.",
|
|
},
|
|
)
|
|
|
|
required.update(
|
|
{
|
|
"diffusion_weight_dtype": (
|
|
["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],
|
|
{
|
|
"default": "default",
|
|
"advanced": True,
|
|
"tooltip": "Weight dtype used only for standalone diffusion models.",
|
|
},
|
|
),
|
|
"grid_columns": ("INT", {"default": 3, "min": 1, "max": 20}),
|
|
"cell_scale": ("FLOAT", {"default": 1.0, "min": 0.05, "max": 4.0, "step": 0.05}),
|
|
"font_size": ("INT", {"default": 22, "min": 8, "max": 96}),
|
|
"padding": ("INT", {"default": 10, "min": 0, "max": 96}),
|
|
"background": (["white", "gray", "black"], {"default": "white"}),
|
|
"show_grid_lines": ("BOOLEAN", {"default": True}),
|
|
"show_cell_labels": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Show the model source and filename above every grid cell.",
|
|
},
|
|
),
|
|
"show_run_header": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Show sampler, scheduler, steps, CFG, and denoise above the grid.",
|
|
},
|
|
),
|
|
"continue_on_error": ("BOOLEAN", {"default": True}),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"required": required,
|
|
"optional": {
|
|
"clip": (
|
|
"CLIP",
|
|
{
|
|
"tooltip": "Required for standalone diffusion models or checkpoints without an embedded CLIP.",
|
|
},
|
|
),
|
|
"vae": (
|
|
"VAE",
|
|
{
|
|
"tooltip": "Required for standalone diffusion models or checkpoints without an embedded VAE.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("grid_image",)
|
|
FUNCTION = "generate_model_matrix"
|
|
CATEGORY = "ComfyUI-KSampler-Matrix-Lab"
|
|
DESCRIPTION = "Compares selected checkpoints and diffusion models with one shared sampler and scheduler, then returns a labeled image grid."
|
|
|
|
def generate_model_matrix(
|
|
self,
|
|
latent_image,
|
|
positive_text,
|
|
negative_text,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
denoise,
|
|
sampler_name,
|
|
scheduler,
|
|
diffusion_weight_dtype,
|
|
grid_columns,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
continue_on_error,
|
|
clip=None,
|
|
vae=None,
|
|
**model_slots,
|
|
):
|
|
selected_models = collect_selected_slots(
|
|
model_slots.get(f"model_{index:02d}")
|
|
for index in range(1, MODEL_SLOT_COUNT + 1)
|
|
)
|
|
validate_selected_names(
|
|
selected_models,
|
|
get_available_model_choices(),
|
|
"model",
|
|
)
|
|
|
|
progress = comfy.utils.ProgressBar(len(selected_models))
|
|
entries = []
|
|
|
|
for selection in selected_models:
|
|
loaded_model = None
|
|
loaded_clip = None
|
|
loaded_vae = None
|
|
effective_clip = None
|
|
effective_vae = None
|
|
positive = None
|
|
negative = None
|
|
sampled_latent = None
|
|
decoded_images = None
|
|
source_type, model_name = parse_model_selection(selection)
|
|
source_label = "Checkpoint" if source_type == "checkpoint" else "Diffusion"
|
|
display_label = f"{source_label}: {model_name}"
|
|
|
|
try:
|
|
(
|
|
loaded_model,
|
|
loaded_clip,
|
|
loaded_vae,
|
|
_,
|
|
_,
|
|
) = load_model_selection(selection, diffusion_weight_dtype)
|
|
|
|
effective_clip = select_clip(loaded_clip, clip)
|
|
effective_vae = select_vae(loaded_vae, vae)
|
|
if not is_valid_vae(effective_vae):
|
|
raise RuntimeError(
|
|
"No compatible VAE is available. Connect the optional VAE input "
|
|
"when a selected checkpoint does not contain a valid VAE or when "
|
|
"comparing standalone diffusion models."
|
|
)
|
|
|
|
positive = encode_prompt_text(effective_clip, positive_text)
|
|
negative = encode_prompt_text(effective_clip, negative_text)
|
|
sampled_latent = run_single_sample(
|
|
loaded_model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise,
|
|
)
|
|
decoded_images = decode_latent_to_image(effective_vae, sampled_latent)
|
|
entries.append(
|
|
{
|
|
"label": display_label,
|
|
"image": tensor_to_pil(decoded_images),
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
if not continue_on_error:
|
|
raise RuntimeError(
|
|
f"Model Matrix Lab failed for '{display_label}': {exc}"
|
|
) from exc
|
|
entries.append(
|
|
{
|
|
"label": display_label,
|
|
"image": None,
|
|
"error": str(exc),
|
|
}
|
|
)
|
|
finally:
|
|
loaded_model = None
|
|
loaded_clip = None
|
|
loaded_vae = None
|
|
effective_clip = None
|
|
effective_vae = None
|
|
positive = None
|
|
negative = None
|
|
sampled_latent = None
|
|
decoded_images = None
|
|
progress.update(1)
|
|
gc.collect()
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
grid = compose_model_grid(
|
|
entries,
|
|
latent_image,
|
|
sampler_name,
|
|
scheduler,
|
|
steps,
|
|
cfg,
|
|
denoise,
|
|
grid_columns,
|
|
cell_scale,
|
|
font_size,
|
|
padding,
|
|
background,
|
|
show_grid_lines,
|
|
show_cell_labels,
|
|
show_run_header,
|
|
)
|
|
return (pil_to_tensor(grid),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KSamplerMatrixLab": KSamplerMatrixLab,
|
|
"ModelMatrixLab": ModelMatrixLab,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KSamplerMatrixLab": "KSampler Matrix Lab",
|
|
"ModelMatrixLab": "Model Matrix Lab",
|
|
}
|